diff --git a/.release-intents/20260818-ragas-replicate-restate-otel2.json b/.release-intents/20260818-ragas-replicate-restate-otel2.json new file mode 100644 index 00000000..8c5fe5ac --- /dev/null +++ b/.release-intents/20260818-ragas-replicate-restate-otel2.json @@ -0,0 +1,8 @@ +{ + "summary": "Repair Ragas, Replicate, and Restate OTel 2.x instrumentation", + "packages": { + "respan-instrumentation-ragas": "patch", + "respan-instrumentation-replicate": "patch", + "respan-instrumentation-restate": "patch" + } +} diff --git a/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_instrumentation.py index fcca56b0..77cb7861 100644 --- a/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_instrumentation.py @@ -5,21 +5,29 @@ import contextvars import functools import importlib +import importlib.metadata import inspect import logging import threading +from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Callable +from typing import Any from opentelemetry import trace from opentelemetry.semconv_ai import SpanAttributes from opentelemetry.trace import Status, StatusCode - -from respan_instrumentation_ragas._serialization import json_string from respan_sdk.constants import ERROR_MESSAGE_ATTR from respan_sdk.constants.llm_logging import LogMethodChoices from respan_sdk.constants.span_attributes import RESPAN_LOG_METHOD, RESPAN_LOG_TYPE from respan_tracing.core.tracer import RespanTracer +from respan_tracing.utils.span_factory import read_propagated_attributes + +from respan_instrumentation_ragas._serialization import ( + exception_message, + json_string, + safe_text, + sensitive_key, +) logger = logging.getLogger(__name__) @@ -27,6 +35,7 @@ _LOCK = threading.RLock() _REFCOUNT = 0 _CAPTURE_CONTENT = True +_ENABLED = False _EVALUATION_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( "respan_ragas_evaluation_depth", default=0 ) @@ -56,16 +65,28 @@ def _is_respan_tracing_enabled() -> bool: def _entity(kind: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str: if kind.startswith("metric"): metric = args[0] if args else None - name = str(getattr(metric, "name", None) or metric.__class__.__name__) + try: + configured_name = getattr(metric, "name", None) + except Exception: # noqa: BLE001 - vendor objects are untrusted. + configured_name = None + name = safe_text(configured_name) or type(metric).__name__ return f"{name}.batch" if kind == "metric_batch" else name if kind == "experiment_row": wrapper = args[0] if args else None - return f"ragas.experiment.row.{getattr(wrapper, '__name__', 'item')}" + try: + name = getattr(wrapper, "__name__", "item") + except Exception: # noqa: BLE001 + name = "item" + return f"ragas.experiment.row.{safe_text(name) or 'item'}" if kind == "experiment_run": wrapper = args[0] if args else None name = kwargs.get("name") - return f"ragas.experiment.{name or getattr(wrapper, '__name__', 'run')}" - return str(kwargs.get("experiment_name") or "ragas.evaluate") + try: + wrapper_name = getattr(wrapper, "__name__", "run") + except Exception: # noqa: BLE001 + wrapper_name = "run" + return f"ragas.experiment.{safe_text(name or wrapper_name) or 'run'}" + return safe_text(kwargs.get("experiment_name")) or "ragas.evaluate" def _input(kind: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: @@ -84,27 +105,46 @@ def _prepare_span( kind: str, args: tuple[Any, ...], kwargs: dict[str, Any], + has_parent: bool, ) -> None: span.set_attribute(RESPAN_LOG_METHOD, LogMethodChoices.TRACING_INTEGRATION.value) span.set_attribute(RESPAN_LOG_TYPE, "task") span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name) - span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_PATH, entity_name) + span.set_attribute( + SpanAttributes.TRACELOOP_ENTITY_PATH, + entity_name if has_parent else "", + ) if _CAPTURE_CONTENT: span.set_attribute( SpanAttributes.TRACELOOP_ENTITY_INPUT, json_string(_input(kind, args, kwargs)), ) + for key, value in read_propagated_attributes().items(): + if key.startswith("respan."): + if sensitive_key(key): + serialized = "[REDACTED]" + elif isinstance(value, str): + serialized = safe_text(value) + elif isinstance(value, bool | int | float): + serialized = value + else: + serialized = json_string(value) + span.set_attribute(key, serialized) def _status_code(value: Any, *, default: int) -> int: - candidates = [value, getattr(value, "response", None)] + try: + response = getattr(value, "response", None) + except Exception: # noqa: BLE001 - exception objects are untrusted. + response = None + candidates = [value, response] for candidate in candidates: for name in ("status_code", "status"): try: code = getattr(candidate, name, None) if isinstance(code, int): return code - except Exception: + except Exception: # noqa: BLE001,S112 - vendor status properties are untrusted. continue return default @@ -119,20 +159,43 @@ def _mark_error(span: Any, exc: BaseException) -> None: status_code = _status_code(exc, default=500) if status_code < 400: status_code = 500 - message = str(exc) + message = exception_message(exc) span.set_attribute("status_code", status_code) span.set_attribute(ERROR_MESSAGE_ATTR, message) - span.set_attribute( - SpanAttributes.TRACELOOP_ENTITY_OUTPUT, - json_string( - { - "status": "error", - "error": type(exc).__name__, - "message": message, - } - ), - ) + if _CAPTURE_CONTENT: + span.set_attribute( + SpanAttributes.TRACELOOP_ENTITY_OUTPUT, + json_string( + { + "status": "error", + "error": type(exc).__name__, + "message": message, + } + ), + ) span.set_status(Status(StatusCode.ERROR, message)) + span.add_event( + "exception", + { + "exception.type": f"{type(exc).__module__}.{type(exc).__name__}", + "exception.message": message, + }, + ) + + +def _has_active_parent() -> bool: + try: + return bool(trace.get_current_span().get_span_context().is_valid) + except Exception: # noqa: BLE001 - non-recording contexts are root candidates. + return False + + +def _tracer() -> Any: + try: + version = importlib.metadata.version("respan-instrumentation-ragas") + except importlib.metadata.PackageNotFoundError: + version = None + return trace.get_tracer(RAGAS_INSTRUMENTATION_NAME, version) def _depth_for_kind(kind: str) -> contextvars.ContextVar[int] | None: @@ -146,20 +209,28 @@ def _depth_for_kind(kind: str) -> contextvars.ContextVar[int] | None: def _sync_wrapper(original: Callable[..., Any], *, kind: str) -> Callable[..., Any]: @functools.wraps(original) def wrapper(*args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return original(*args, **kwargs) depth = _depth_for_kind(kind) if depth is not None and depth.get() > 0: return original(*args, **kwargs) token = depth.set(depth.get() + 1) if depth is not None else None entity_name = _entity(kind, args, kwargs) + has_parent = _has_active_parent() try: - tracer = trace.get_tracer(RAGAS_INSTRUMENTATION_NAME) - with tracer.start_as_current_span(f"{entity_name}.task") as span: + tracer = _tracer() + with tracer.start_as_current_span( + f"{entity_name}.task", + record_exception=False, + set_status_on_exception=False, + ) as span: _prepare_span( span, entity_name=entity_name, kind=kind, args=args, kwargs=kwargs, + has_parent=has_parent, ) try: result = original(*args, **kwargs) @@ -179,20 +250,28 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: def _async_wrapper(original: Callable[..., Any], *, kind: str) -> Callable[..., Any]: @functools.wraps(original) async def wrapper(*args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return await original(*args, **kwargs) depth = _depth_for_kind(kind) if depth is not None and depth.get() > 0: return await original(*args, **kwargs) token = depth.set(depth.get() + 1) if depth is not None else None entity_name = _entity(kind, args, kwargs) + has_parent = _has_active_parent() try: - tracer = trace.get_tracer(RAGAS_INSTRUMENTATION_NAME) - with tracer.start_as_current_span(f"{entity_name}.task") as span: + tracer = _tracer() + with tracer.start_as_current_span( + f"{entity_name}.task", + record_exception=False, + set_status_on_exception=False, + ) as span: _prepare_span( span, entity_name=entity_name, kind=kind, args=args, kwargs=kwargs, + has_parent=has_parent, ) try: result = await original(*args, **kwargs) @@ -241,15 +320,15 @@ def _install_patches() -> None: _patch(owner, "evaluate", kind="evaluation") _patch(owner, "aevaluate", kind="evaluation") - single_turn = getattr(metrics_base, "SingleTurnMetric") + single_turn = metrics_base.SingleTurnMetric _patch(single_turn, "single_turn_score", kind="metric") _patch(single_turn, "single_turn_ascore", kind="metric") - multi_turn = getattr(metrics_base, "MultiTurnMetric") + multi_turn = metrics_base.MultiTurnMetric _patch(multi_turn, "multi_turn_score", kind="metric") _patch(multi_turn, "multi_turn_ascore", kind="metric") importlib.import_module("ragas.metrics.collections") collections_base = importlib.import_module("ragas.metrics.collections.base") - collection_metric = getattr(collections_base, "BaseMetric") + collection_metric = collections_base.BaseMetric _patch(collection_metric, "score", kind="metric") _patch(collection_metric, "batch_score", kind="metric_batch") _patch(collection_metric, "abatch_score", kind="metric_batch") @@ -257,7 +336,7 @@ def _install_patches() -> None: if "ascore" in metric_class.__dict__: _patch(metric_class, "ascore", kind="metric") - wrapper = getattr(experiment, "ExperimentWrapper") + wrapper = experiment.ExperimentWrapper _patch(wrapper, "__call__", kind="experiment_row") _patch(wrapper, "arun", kind="experiment_run") @@ -279,7 +358,7 @@ def __init__(self, *, capture_content: bool = True) -> None: self._is_instrumented = False def activate(self) -> None: - global _CAPTURE_CONTENT, _REFCOUNT + global _CAPTURE_CONTENT, _ENABLED, _REFCOUNT if self._is_instrumented or not _is_respan_tracing_enabled(): return @@ -292,16 +371,22 @@ def activate(self) -> None: with _LOCK: if _REFCOUNT == 0: _CAPTURE_CONTENT = self._capture_content - _install_patches() + try: + _install_patches() + except Exception: + _remove_patches() + raise + _ENABLED = True elif _CAPTURE_CONTENT != self._capture_content: - logger.warning( - "Ragas is already instrumented; the first capture_content setting wins" + raise ValueError( + "all active RagasInstrumentor instances must use the same " + "capture_content setting" ) _REFCOUNT += 1 self._is_instrumented = True def deactivate(self) -> None: - global _REFCOUNT + global _ENABLED, _REFCOUNT if not self._is_instrumented: return @@ -309,4 +394,5 @@ def deactivate(self) -> None: self._is_instrumented = False _REFCOUNT = max(0, _REFCOUNT - 1) if _REFCOUNT == 0: + _ENABLED = False _remove_patches() diff --git a/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_serialization.py b/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_serialization.py index 76406506..9a8ecfa1 100644 --- a/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_serialization.py +++ b/python-sdks/instrumentations/respan-instrumentation-ragas/src/respan_instrumentation_ragas/_serialization.py @@ -1,64 +1,198 @@ -"""Bounded JSON serialization for Ragas values.""" +"""Bounded, privacy-safe JSON serialization for Ragas values.""" from __future__ import annotations -import dataclasses import json +import math +import re from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass +from enum import Enum +from itertools import islice +from numbers import Integral, Real from typing import Any -_MAX_DEPTH = 5 -_MAX_ITEMS = 25 -_MAX_STRING = 8_000 +MAX_ATTRIBUTE_BYTES = 16_000 +MAX_DEPTH = 8 +MAX_ITEMS = 50 +MAX_STRING_BYTES = 4_000 +REDACTED = "[REDACTED]" +_SENSITIVE_SUFFIXES = ( + "apikey", + "authorization", + "credential", + "password", + "secret", + "sessiontoken", + "token", +) +_SECRET_ASSIGNMENT = re.compile( + r"(?i)([\"']?(?:api[_-]?key|authorization|password|secret|session[_-]?token|token)[\"']?)" + r"(\s*[:=]\s*)([\"']?)([^\s,;}\"']+)([\"']?)" +) +_BEARER = re.compile(r"(?i)\bbearer\s+[a-z0-9._~+/=-]+") +_ADDRESS = re.compile(r"(?i)\b0x[0-9a-f]{6,}\b") -def _value(value: Any, *, depth: int, seen: set[int]) -> Any: - if value is None or isinstance(value, bool | int | float): + +def _truncate_utf8(value: str, limit: int = MAX_STRING_BYTES) -> str: + encoded = value.encode("utf-8") + if len(encoded) <= limit: return value + suffix = "...[truncated]" + budget = max(0, limit - len(suffix.encode("utf-8"))) + return encoded[:budget].decode("utf-8", errors="ignore") + suffix + + +def safe_text(value: Any, *, default: str = "") -> str: if isinstance(value, str): - return value[:_MAX_STRING] - if depth >= _MAX_DEPTH: - return repr(value)[:_MAX_STRING] + value = _BEARER.sub(REDACTED, value) + value = _SECRET_ASSIGNMENT.sub( + lambda match: ( + f"{match.group(1)}{match.group(2)}{match.group(3)}" + f"{REDACTED}{match.group(5)}" + ), + value, + ) + return _truncate_utf8(_ADDRESS.sub("0x", value)) + if value is None: + return default + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, Integral): + return str(int(value)) + if isinstance(value, Real) and math.isfinite(float(value)): + return str(float(value)) + return f"<{type(value).__name__}>" - value_id = id(value) - if value_id in seen: - return "" - seen.add(value_id) + +def exception_message(exc: BaseException) -> str: try: - if dataclasses.is_dataclass(value) and not isinstance(value, type): - value = dataclasses.asdict(value) - elif callable(getattr(value, "model_dump", None)): - value = value.model_dump() - elif callable(getattr(value, "to_dict", None)): - value = value.to_dict() - elif callable(getattr(value, "dict", None)): - value = value.dict() + arguments = exc.args + except Exception: # noqa: BLE001 - hostile exception objects are supported. + arguments = () + for argument in arguments: + if isinstance(argument, str | bool | int | float): + return safe_text(argument) + return type(exc).__name__ - if isinstance(value, Mapping): - return { - str(key): _value(item, depth=depth + 1, seen=seen) - for key, item in list(value.items())[:_MAX_ITEMS] + +def _key(value: Any) -> str: + if isinstance(value, Enum): + value = value.value + return safe_text(value)[:256] + + +def sensitive_key(value: Any) -> bool: + if not isinstance(value, str): + return False + normalized = re.sub(r"[^a-z0-9]", "", value.lower()) + return any(normalized.endswith(suffix) for suffix in _SENSITIVE_SUFFIXES) + + +def json_value(value: Any, *, depth: int = 0, seen: set[int] | None = None) -> Any: + if value is None or isinstance(value, bool): + return value + if isinstance(value, str): + return safe_text(value) + if isinstance(value, Integral): + return int(value) + if isinstance(value, Real): + number = float(value) + return number if math.isfinite(number) else None + if isinstance(value, bytes | bytearray | memoryview): + return {"length": len(value), "type": type(value).__name__} + if isinstance(value, Enum): + return json_value(value.value, depth=depth + 1, seen=seen) + if depth >= MAX_DEPTH: + return {"truncated": "max_depth", "type": type(value).__name__} + + active = seen if seen is not None else set() + identity = id(value) + if identity in active: + return "" + active.add(identity) + try: + if is_dataclass(value) and not isinstance(value, type): + selected = list(islice(fields(value), MAX_ITEMS + 1)) + converted = { + field.name: getattr(value, field.name) for field in selected[:MAX_ITEMS] } - if isinstance(value, Sequence) and not isinstance(value, str | bytes): - return [ - _value(item, depth=depth + 1, seen=seen) - for item in list(value)[:_MAX_ITEMS] + if len(selected) > MAX_ITEMS: + converted["__truncated_items__"] = True + return json_value(converted, depth=depth + 1, seen=active) + if isinstance(value, Mapping): + result: dict[str, Any] = {} + items = list(islice(value.items(), MAX_ITEMS + 1)) + for key, item in items[:MAX_ITEMS]: + key_text = _key(key) + result[key_text] = ( + REDACTED + if sensitive_key(key) + else json_value(item, depth=depth + 1, seen=active) + ) + if len(items) > MAX_ITEMS: + result["__truncated_items__"] = True + return result + if isinstance(value, Sequence) and not isinstance( + value, str | bytes | bytearray + ): + items = list(islice(iter(value), MAX_ITEMS + 1)) + converted = [ + json_value(item, depth=depth + 1, seen=active) + for item in items[:MAX_ITEMS] ] - if hasattr(value, "__dict__"): - public = { - key: item - for key, item in vars(value).items() - if not key.startswith("_") - } - if public: - return _value(public, depth=depth + 1, seen=seen) - return repr(value)[:_MAX_STRING] - except Exception: - return repr(value)[:_MAX_STRING] + if len(items) > MAX_ITEMS: + return {"items": converted, "truncated": True} + return converted + for method_name in ("model_dump", "to_dict", "dict"): + try: + method = getattr(value, method_name, None) + except Exception: # noqa: BLE001 - vendor attributes are untrusted. + method = None + if not callable(method): + continue + try: + converted = method() + except Exception: # noqa: BLE001,S112 - fall through to a type summary. + continue + if isinstance(converted, Mapping): + return json_value(converted, depth=depth + 1, seen=active) + return {"type": type(value).__name__} + except Exception: # noqa: BLE001 - serialization must never break Ragas. + return {"type": type(value).__name__, "unserializable": True} finally: - seen.discard(value_id) + active.discard(identity) def json_string(value: Any) -> str: - """Return a deterministic, bounded JSON representation.""" - return json.dumps(_value(value, depth=0, seen=set()), default=str, sort_keys=True) + """Return valid JSON bounded by the OTel attribute budget.""" + encoded = json.dumps( + json_value(value), + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + encoded_bytes = len(encoded.encode("utf-8")) + if encoded_bytes <= MAX_ATTRIBUTE_BYTES: + return encoded + low, high, result = 0, len(encoded), "" + while low <= high: + midpoint = (low + high) // 2 + candidate = json.dumps( + { + "original_bytes": encoded_bytes, + "preview": encoded[:midpoint], + "truncated": True, + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + if len(candidate.encode("utf-8")) <= MAX_ATTRIBUTE_BYTES: + result = candidate + low = midpoint + 1 + else: + high = midpoint - 1 + return result diff --git a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_backend_status_and_lifecycle.py b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_backend_status_and_lifecycle.py index 8d96c2a7..c6bf64f1 100644 --- a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_backend_status_and_lifecycle.py +++ b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_backend_status_and_lifecycle.py @@ -3,12 +3,11 @@ import json import pytest +import respan_instrumentation_ragas._instrumentation as instrumentation from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.semconv_ai import SpanAttributes - -import respan_instrumentation_ragas._instrumentation as instrumentation from respan_instrumentation_ragas import RagasInstrumentor @@ -21,6 +20,8 @@ def _exporter(monkeypatch) -> InMemorySpanExporter: "get_tracer", lambda *args, **kwargs: provider.get_tracer("test.ragas.status"), ) + monkeypatch.setattr(instrumentation, "_ENABLED", True) + monkeypatch.setattr(instrumentation, "_CAPTURE_CONTENT", True) return exporter diff --git a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_modern_collections.py b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_modern_collections.py index a20a15ef..aed08850 100644 --- a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_modern_collections.py +++ b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_modern_collections.py @@ -2,14 +2,14 @@ import asyncio import json + import pytest +import respan_instrumentation_ragas._instrumentation as instrumentation from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.semconv_ai import SpanAttributes from ragas.metrics.collections import ExactMatch, StringPresence - -import respan_instrumentation_ragas._instrumentation as instrumentation from respan_instrumentation_ragas import RagasInstrumentor diff --git a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_ragas_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_ragas_instrumentation.py index 20809636..b2337dde 100644 --- a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_ragas_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_ragas_instrumentation.py @@ -5,13 +5,12 @@ from types import SimpleNamespace import pytest +import respan_instrumentation_ragas._instrumentation as instrumentation from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.semconv_ai import SpanAttributes from opentelemetry.trace import StatusCode - -import respan_instrumentation_ragas._instrumentation as instrumentation from respan_sdk.constants.span_attributes import RESPAN_LOG_METHOD, RESPAN_LOG_TYPE @@ -26,6 +25,7 @@ def spans(monkeypatch): lambda *args, **kwargs: provider.get_tracer("test.ragas"), ) monkeypatch.setattr(instrumentation, "_CAPTURE_CONTENT", True) + monkeypatch.setattr(instrumentation, "_ENABLED", True) return exporter diff --git a/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_real_ragas_runtime.py b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_real_ragas_runtime.py new file mode 100644 index 00000000..277317e0 --- /dev/null +++ b/python-sdks/instrumentations/respan-instrumentation-ragas/tests/test_real_ragas_runtime.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import asyncio +import json +from collections import Counter + +import ragas +import respan_instrumentation_ragas._instrumentation as instrumentation +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.semconv_ai import SpanAttributes +from ragas import EvaluationDataset +from ragas.backends.inmemory import InMemoryBackend +from ragas.dataset import Dataset +from ragas.metrics import ExactMatch +from respan_instrumentation_ragas import RagasInstrumentor +from respan_instrumentation_ragas._serialization import json_string +from respan_sdk.constants.span_attributes import RESPAN_LOG_TYPE + + +def test_real_current_evaluate_and_experiment_export_connected_spans( + monkeypatch, +) -> None: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr( + instrumentation.trace, + "get_tracer", + lambda *args, **kwargs: provider.get_tracer("ragas", "0.1.0"), + ) + monkeypatch.setattr(instrumentation, "_REFCOUNT", 0) + monkeypatch.setattr(instrumentation, "_PATCHES", []) + monkeypatch.setattr(instrumentation, "_ENABLED", False) + + adapter = RagasInstrumentor() + adapter.activate() + evaluation_dataset = EvaluationDataset.from_list( + [{"user_input": "Capital?", "response": "Paris", "reference": "Paris"}] + ) + result = ragas.evaluate( + evaluation_dataset, + metrics=[ExactMatch()], + show_progress=False, + ) + assert result["exact_match"] == [1.0] + + backend = InMemoryBackend() + + @ragas.experiment(backend=backend, name_prefix="offline") + async def answer(row): + return {"answer": row["answer"]} + + dataset = Dataset( + name="questions", + backend=backend, + data=[{"answer": "Paris"}, {"answer": "Rome"}], + ) + experiment = asyncio.run(answer.arun(dataset, name="two-rows")) + assert len(experiment) == 2 + adapter.deactivate() + + spans = exporter.get_finished_spans() + names = Counter( + span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] for span in spans + ) + assert names == Counter( + { + "exact_match": 1, + "ragas.evaluate": 1, + "ragas.experiment.two-rows": 1, + "ragas.experiment.row.answer": 2, + } + ) + assert len({span.context.span_id for span in spans}) == 5 + assert all(span.attributes[RESPAN_LOG_TYPE] == "task" for span in spans) + assert all("traceloop.span.kind" not in span.attributes for span in spans) + roots = [span for span in spans if span.parent is None] + assert len(roots) == 2 + assert all( + root.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] == "" for root in roots + ) + parents = {span.context.span_id for span in spans} + assert all(span.parent is None or span.parent.span_id in parents for span in spans) + assert all(span.instrumentation_scope.name == "ragas" for span in spans) + + +def test_ragas_serialization_is_valid_bounded_and_private() -> None: + class Hostile: + def __str__(self) -> str: + raise AssertionError("must not stringify") + + def __repr__(self) -> str: + raise AssertionError("must not repr") + + encoded = json_string( + { + "api_token": "plain-secret", + "content": "😀" * 5_000, + "hostile": Hostile(), + "nonfinite": float("inf"), + } + ) + assert len(encoded.encode("utf-8")) <= 16_000 + assert "plain-secret" not in encoded + assert json.loads(encoded)["nonfinite"] is None diff --git a/python-sdks/instrumentations/respan-instrumentation-replicate/README.md b/python-sdks/instrumentations/respan-instrumentation-replicate/README.md index 89cecb65..329705cc 100644 --- a/python-sdks/instrumentations/respan-instrumentation-replicate/README.md +++ b/python-sdks/instrumentations/respan-instrumentation-replicate/README.md @@ -33,7 +33,11 @@ def run_prediction() -> str: "meta/meta-llama-3-8b-instruct", input={"prompt": "Reply with one concise sentence about tracing."}, ) - return "".join(str(chunk) for chunk in output) if not isinstance(output, str) else output + return ( + "".join(str(chunk) for chunk in output) + if not isinstance(output, str) + else output + ) print(run_prediction()) diff --git a/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_instrumentation.py index 298008f7..df74ed5a 100644 --- a/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_instrumentation.py @@ -4,13 +4,27 @@ import contextlib import contextvars +import functools import importlib +import importlib.metadata import logging +import threading import time -from collections.abc import AsyncIterator, Iterator -from typing import Any, Callable +from collections.abc import AsyncIterator, Callable, Iterator +from dataclasses import dataclass +from types import TracebackType +from typing import Any, Self from opentelemetry import trace +from opentelemetry.sdk.util.instrumentation import InstrumentationScope +from opentelemetry.semconv_ai import SpanAttributes +from respan_sdk.constants import ERROR_MESSAGE_ATTR +from respan_sdk.utils.data_processing.id_processing import ( + format_span_id, + format_trace_id, +) +from respan_tracing.core.tracer import RespanTracer +from respan_tracing.utils.span_factory import build_readable_span, inject_span from respan_instrumentation_replicate._constants import ( ASYNC_PREFIX, @@ -25,20 +39,34 @@ RESPAN_PARAMS_KEY, RESPAN_PARAMS_MODEL_KEY, ) +from respan_instrumentation_replicate._serialization import ( + exception_message, + exception_status, + prediction_summary, + safe_text, +) from respan_instrumentation_replicate._translator import ( build_model_call_span_data, build_operation_span_data, - output_to_text, ) -from respan_sdk.utils.data_processing.id_processing import ( - format_span_id, - format_trace_id, -) -from respan_tracing.core.tracer import RespanTracer -from respan_tracing.utils.span_factory import build_readable_span, inject_span logger = logging.getLogger(__name__) +_LOCK = threading.RLock() +_REFCOUNT = 0 +_ENABLED = False + + +@dataclass +class _Patch: + owner: Any + name: str + original: Any + replacement: Any + + +_PATCHES: list[_Patch] = [] + _SUPPRESSED_SPAN_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( "respan_replicate_suppressed_span_depth", default=0, @@ -62,7 +90,7 @@ def _current_otel_parent() -> tuple[str | None, str | None]: current_span = trace.get_current_span() try: span_context = current_span.get_span_context() - except Exception: + except Exception: # noqa: BLE001 - non-recording spans can be hostile proxies. return None, None trace_id = getattr(span_context, "trace_id", 0) @@ -80,13 +108,26 @@ def _emit_span( attributes: dict[str, Any], start_time_ns: int, end_time_ns: int | None = None, - error: Exception | None = None, + error: BaseException | None = None, parent_context: tuple[str | None, str | None] | None = None, ) -> None: if _spans_suppressed(): return trace_id, parent_id = parent_context or _current_otel_parent() + attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] = ( + attributes.get(SpanAttributes.TRACELOOP_ENTITY_PATH, span_name) + if parent_id + else "" + ) + status_code = exception_status(error) if error is not None else 200 + if error is not None: + message = exception_message(error) + attributes["status_code"] = status_code + attributes[ERROR_MESSAGE_ATTR] = message + else: + message = None + attributes.setdefault("status_code", 200) span = build_readable_span( name=span_name, trace_id=trace_id, @@ -94,8 +135,16 @@ def _emit_span( start_time_ns=start_time_ns, end_time_ns=end_time_ns or time.time_ns(), attributes=attributes, - status_code=500 if error is not None else 200, - error_message=str(error) if error is not None else None, + status_code=status_code, + error_message=message, + ) + try: + package_version = importlib.metadata.version("respan-instrumentation-replicate") + except importlib.metadata.PackageNotFoundError: + package_version = None + span._instrumentation_scope = InstrumentationScope( # type: ignore[attr-defined] + REPLICATE_INSTRUMENTATION_NAME, + package_version, ) inject_span(span=span) @@ -110,7 +159,7 @@ def _reported_model_from_respan_params(respan_params: Any) -> str | None: if not isinstance(respan_params, dict): return None model = respan_params.get(RESPAN_PARAMS_MODEL_KEY) - return str(model) if model else None + return safe_text(model) if model else None def _set_prediction_reported_model(prediction: Any, respan_params: Any) -> None: @@ -119,10 +168,10 @@ def _set_prediction_reported_model(prediction: Any, respan_params: Any) -> None: return try: object.__setattr__(prediction, PREDICTION_RESPAN_MODEL_ATTR, reported_model) - except Exception: + except Exception: # noqa: BLE001 - resources may reject private attributes. try: setattr(prediction, PREDICTION_RESPAN_MODEL_ATTR, reported_model) - except Exception: + except Exception: # noqa: BLE001 - best-effort metadata only. return @@ -143,14 +192,14 @@ def __init__( self, *, iterator: Iterator[Any], - emit_once: Callable[[list[Any], Exception | None], None], + emit_once: Callable[[list[Any], BaseException | None], None], ) -> None: self._iterator = iterator self._emit_once = emit_once self._chunks: list[Any] = [] self._emitted = False - def __iter__(self) -> "_SyncIteratorProxy": + def __iter__(self) -> _SyncIteratorProxy: return self def __next__(self) -> Any: @@ -160,7 +209,7 @@ def __next__(self) -> Any: except StopIteration: self._emit(error=None) raise - except Exception as exc: + except BaseException as exc: self._emit(error=exc) raise if len(self._chunks) < MAX_STREAM_CHUNKS: @@ -170,7 +219,40 @@ def __next__(self) -> Any: def __getattr__(self, name: str) -> Any: return getattr(self._iterator, name) - def _emit(self, *, error: Exception | None) -> None: + def close(self) -> None: + error: BaseException | None = None + try: + close = getattr(self._iterator, "close", None) + if callable(close): + close() + except BaseException as exc: + error = exc + raise + finally: + self._emit(error=error) + + def __enter__(self) -> Self: + enter = getattr(self._iterator, "__enter__", None) + if callable(enter): + enter() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> Any: + try: + exit_method = getattr(self._iterator, "__exit__", None) + if callable(exit_method): + return exit_method(exc_type, exc, tb) + self.close() + return None + finally: + self._emit(error=exc) + + def _emit(self, *, error: BaseException | None) -> None: if self._emitted: return self._emitted = True @@ -182,14 +264,14 @@ def __init__( self, *, iterator: AsyncIterator[Any], - emit_once: Callable[[list[Any], Exception | None], None], + emit_once: Callable[[list[Any], BaseException | None], None], ) -> None: self._iterator = iterator self._chunks: list[Any] = [] self._emit_once = emit_once self._emitted = False - def __aiter__(self) -> "_AsyncIteratorProxy": + def __aiter__(self) -> _AsyncIteratorProxy: return self async def __anext__(self) -> Any: @@ -199,7 +281,7 @@ async def __anext__(self) -> Any: except StopAsyncIteration: self._emit(error=None) raise - except Exception as exc: + except BaseException as exc: self._emit(error=exc) raise if len(self._chunks) < MAX_STREAM_CHUNKS: @@ -209,7 +291,46 @@ async def __anext__(self) -> Any: def __getattr__(self, name: str) -> Any: return getattr(self._iterator, name) - def _emit(self, *, error: Exception | None) -> None: + async def aclose(self) -> None: + error: BaseException | None = None + try: + close = getattr(self._iterator, "aclose", None) + if callable(close): + await close() + else: + close = getattr(self._iterator, "close", None) + if callable(close): + result = close() + if hasattr(result, "__await__"): + await result + except BaseException as exc: + error = exc + raise + finally: + self._emit(error=error) + + async def __aenter__(self) -> Self: + enter = getattr(self._iterator, "__aenter__", None) + if callable(enter): + await enter() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> Any: + try: + exit_method = getattr(self._iterator, "__aexit__", None) + if callable(exit_method): + return await exit_method(exc_type, exc, tb) + await self.aclose() + return None + finally: + self._emit(error=exc) + + def _emit(self, *, error: BaseException | None) -> None: if self._emitted: return self._emitted = True @@ -217,7 +338,14 @@ def _emit(self, *, error: Exception | None) -> None: def _wrap_sync_run(original: Any, *, span_name: str, stream: bool = False) -> Any: - def wrapper(self: Any, ref: Any, input: Any = None, *args: Any, **kwargs: Any) -> Any: + @functools.wraps(original) + def wrapper( + self: Any, ref: Any, input: Any = None, *args: Any, **kwargs: Any + ) -> Any: + if not _ENABLED: + if stream: + return original(self, ref, *args, input=input, **kwargs) + return original(self, ref, input, *args, **kwargs) call_kwargs, respan_params = _pop_respan_params(kwargs) event_kwargs = {**call_kwargs, RESPAN_PARAMS_KEY: respan_params} parent_context = _current_otel_parent() @@ -228,7 +356,7 @@ def wrapper(self: Any, ref: Any, input: Any = None, *args: Any, **kwargs: Any) - output = original(self, ref, *args, input=input, **call_kwargs) else: output = original(self, ref, input, *args, **call_kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_model_call_span_data( span_name=span_name, ref=ref, @@ -248,7 +376,7 @@ def wrapper(self: Any, ref: Any, input: Any = None, *args: Any, **kwargs: Any) - if _is_sync_iterator(output): - def emit_once(chunks: list[Any], error: Exception | None) -> None: + def emit_once(chunks: list[Any], error: BaseException | None) -> None: resolved_span_name, attrs = build_model_call_span_data( span_name=span_name, ref=ref, @@ -288,9 +416,14 @@ def emit_once(chunks: list[Any], error: Exception | None) -> None: def _wrap_async_run(original: Any, *, span_name: str, stream: bool = False) -> Any: + @functools.wraps(original) async def wrapper( self: Any, ref: Any, input: Any = None, *args: Any, **kwargs: Any ) -> Any: + if not _ENABLED: + if stream: + return await original(self, ref, input=input, **kwargs) + return await original(self, ref, input, *args, **kwargs) call_kwargs, respan_params = _pop_respan_params(kwargs) event_kwargs = {**call_kwargs, RESPAN_PARAMS_KEY: respan_params} parent_context = _current_otel_parent() @@ -301,7 +434,7 @@ async def wrapper( output = await original(self, ref, input=input, **call_kwargs) else: output = await original(self, ref, input, *args, **call_kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_model_call_span_data( span_name=span_name, ref=ref, @@ -321,7 +454,7 @@ async def wrapper( if _is_async_iterator(output): - def emit_once(chunks: list[Any], error: Exception | None) -> None: + def emit_once(chunks: list[Any], error: BaseException | None) -> None: resolved_span_name, attrs = build_model_call_span_data( span_name=span_name, ref=ref, @@ -363,7 +496,10 @@ def emit_once(chunks: list[Any], error: Exception | None) -> None: def _wrap_prediction_create(original: Any, *, is_async: bool = False) -> Any: if is_async: + @functools.wraps(original) async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return await original(self, *args, **kwargs) call_kwargs, respan_params = _pop_respan_params(kwargs) event_kwargs = {**call_kwargs, RESPAN_PARAMS_KEY: respan_params} parent_context = _current_otel_parent() @@ -371,7 +507,7 @@ async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: try: with _suppress_nested_spans(): prediction = await original(self, *args, **call_kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_model_call_span_data( span_name=REPLICATE_PREDICTION_CREATE_SPAN_NAME, ref=args[0] if args else None, @@ -406,7 +542,10 @@ async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: return async_wrapper + @functools.wraps(original) def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return original(self, *args, **kwargs) call_kwargs, respan_params = _pop_respan_params(kwargs) event_kwargs = {**call_kwargs, RESPAN_PARAMS_KEY: respan_params} parent_context = _current_otel_parent() @@ -414,7 +553,7 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: try: with _suppress_nested_spans(): prediction = original(self, *args, **call_kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_model_call_span_data( span_name=REPLICATE_PREDICTION_CREATE_SPAN_NAME, ref=args[0] if args else None, @@ -453,16 +592,19 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: def _wrap_prediction_wait(original: Any, *, is_async: bool = False) -> Any: if is_async: + @functools.wraps(original) async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return await original(self, *args, **kwargs) parent_context = _current_otel_parent() start_ns = time.time_ns() try: with _suppress_nested_spans(): result = await original(self, *args, **kwargs) - except Exception as exc: - resolved_span_name, attrs = build_model_call_span_data( + except BaseException as exc: + resolved_span_name, attrs = build_operation_span_data( span_name=REPLICATE_PREDICTION_WAIT_SPAN_NAME, - prediction=self, + input_value=prediction_summary(self), error=exc, ) _emit_span( @@ -474,9 +616,10 @@ async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: ) raise - resolved_span_name, attrs = build_model_call_span_data( + resolved_span_name, attrs = build_operation_span_data( span_name=REPLICATE_PREDICTION_WAIT_SPAN_NAME, - prediction=self, + input_value=prediction_summary(self), + output=prediction_summary(self), ) _emit_span( span_name=resolved_span_name, @@ -488,16 +631,19 @@ async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: return async_wrapper + @functools.wraps(original) def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return original(self, *args, **kwargs) parent_context = _current_otel_parent() start_ns = time.time_ns() try: with _suppress_nested_spans(): result = original(self, *args, **kwargs) - except Exception as exc: - resolved_span_name, attrs = build_model_call_span_data( + except BaseException as exc: + resolved_span_name, attrs = build_operation_span_data( span_name=REPLICATE_PREDICTION_WAIT_SPAN_NAME, - prediction=self, + input_value=prediction_summary(self), error=exc, ) _emit_span( @@ -509,9 +655,10 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: ) raise - resolved_span_name, attrs = build_model_call_span_data( + resolved_span_name, attrs = build_operation_span_data( span_name=REPLICATE_PREDICTION_WAIT_SPAN_NAME, - prediction=self, + input_value=prediction_summary(self), + output=prediction_summary(self), ) _emit_span( span_name=resolved_span_name, @@ -527,13 +674,16 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: def _wrap_prediction_stream(original: Any, *, is_async: bool = False) -> Any: if is_async: + @functools.wraps(original) def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return original(self, *args, **kwargs) parent_context = _current_otel_parent() start_ns = time.time_ns() try: with _suppress_nested_spans(): iterator = original(self, *args, **kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_model_call_span_data( span_name=REPLICATE_STREAM_SPAN_NAME, prediction=self, @@ -549,7 +699,7 @@ def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: ) raise - def emit_once(chunks: list[Any], error: Exception | None) -> None: + def emit_once(chunks: list[Any], error: BaseException | None) -> None: resolved_span_name, attrs = build_model_call_span_data( span_name=REPLICATE_STREAM_SPAN_NAME, prediction=self, @@ -569,13 +719,16 @@ def emit_once(chunks: list[Any], error: Exception | None) -> None: return async_wrapper + @functools.wraps(original) def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return original(self, *args, **kwargs) parent_context = _current_otel_parent() start_ns = time.time_ns() try: with _suppress_nested_spans(): iterator = original(self, *args, **kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_model_call_span_data( span_name=REPLICATE_STREAM_SPAN_NAME, prediction=self, @@ -591,7 +744,7 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: ) raise - def emit_once(chunks: list[Any], error: Exception | None) -> None: + def emit_once(chunks: list[Any], error: BaseException | None) -> None: resolved_span_name, attrs = build_model_call_span_data( span_name=REPLICATE_STREAM_SPAN_NAME, prediction=self, @@ -615,14 +768,17 @@ def emit_once(chunks: list[Any], error: Exception | None) -> None: def _wrap_operation(original: Any, *, span_name: str, is_async: bool = False) -> Any: if is_async: + @functools.wraps(original) async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return await original(self, *args, **kwargs) parent_context = _current_otel_parent() start_ns = time.time_ns() input_value = {"args": args, "kwargs": kwargs} try: with _suppress_nested_spans(): result = await original(self, *args, **kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_operation_span_data( span_name=span_name, input_value=input_value, @@ -640,7 +796,7 @@ async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: resolved_span_name, attrs = build_operation_span_data( span_name=span_name, input_value=input_value, - output=output_to_text(result), + output=result, ) _emit_span( span_name=resolved_span_name, @@ -652,14 +808,17 @@ async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: return async_wrapper + @functools.wraps(original) def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if not _ENABLED: + return original(self, *args, **kwargs) parent_context = _current_otel_parent() start_ns = time.time_ns() input_value = {"args": args, "kwargs": kwargs} try: with _suppress_nested_spans(): result = original(self, *args, **kwargs) - except Exception as exc: + except BaseException as exc: resolved_span_name, attrs = build_operation_span_data( span_name=span_name, input_value=input_value, @@ -677,7 +836,7 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: resolved_span_name, attrs = build_operation_span_data( span_name=span_name, input_value=input_value, - output=output_to_text(result), + output=result, ) _emit_span( span_name=resolved_span_name, @@ -692,7 +851,7 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: def _module_run_wrapper(module: Any, *, method_name: str) -> Callable[..., Any]: def wrapper(ref: Any, input: Any = None, *args: Any, **kwargs: Any) -> Any: - client = getattr(module, "default_client") + client = module.default_client method = getattr(client, method_name) return method(ref, input, *args, **kwargs) @@ -701,7 +860,7 @@ def wrapper(ref: Any, input: Any = None, *args: Any, **kwargs: Any) -> Any: def _module_async_run_wrapper(module: Any, *, method_name: str) -> Callable[..., Any]: async def wrapper(ref: Any, input: Any = None, *args: Any, **kwargs: Any) -> Any: - client = getattr(module, "default_client") + client = module.default_client method = getattr(client, method_name) return await method(ref, input, *args, **kwargs) @@ -710,16 +869,18 @@ async def wrapper(ref: Any, input: Any = None, *args: Any, **kwargs: Any) -> Any def _module_stream_wrapper(module: Any, *, method_name: str) -> Callable[..., Any]: def wrapper(ref: Any, *, input: Any = None, **kwargs: Any) -> Any: - client = getattr(module, "default_client") + client = module.default_client method = getattr(client, method_name) return method(ref, input=input, **kwargs) return wrapper -def _module_async_stream_wrapper(module: Any, *, method_name: str) -> Callable[..., Any]: +def _module_async_stream_wrapper( + module: Any, *, method_name: str +) -> Callable[..., Any]: async def wrapper(ref: Any, input: Any = None, **kwargs: Any) -> Any: - client = getattr(module, "default_client") + client = module.default_client method = getattr(client, method_name) return await method(ref, input=input, **kwargs) @@ -732,7 +893,6 @@ class ReplicateInstrumentor: name = REPLICATE_INSTRUMENTATION_NAME def __init__(self) -> None: - self._originals: dict[tuple[Any, str], Any] = {} self._is_instrumented = False @staticmethod @@ -742,14 +902,13 @@ def _is_respan_tracing_enabled() -> bool: return True return bool(getattr(tracer, "is_enabled", True)) - def _patch_attr(self, owner: Any, attr_name: str, replacement: Any) -> None: - key = (owner, attr_name) - if key not in self._originals: - self._originals[key] = getattr(owner, attr_name) + @staticmethod + def _patch_attr(owner: Any, attr_name: str, replacement: Any) -> None: + original = getattr(owner, attr_name) setattr(owner, attr_name, replacement) + _PATCHES.append(_Patch(owner, attr_name, original, replacement)) - def activate(self) -> None: - """Monkey-patch the Replicate SDK.""" + def _activate_once(self) -> None: if self._is_instrumented: return @@ -770,15 +929,15 @@ def activate(self) -> None: ) return - Client = getattr(client_module, "Client") - Predictions = getattr(prediction_module, "Predictions") - Prediction = getattr(prediction_module, "Prediction") + Client = client_module.Client + Predictions = prediction_module.Predictions + Prediction = prediction_module.Prediction self._patch_attr( Client, "run", _wrap_sync_run( - getattr(Client, "run"), + Client.run, span_name=REPLICATE_RUN_SPAN_NAME, ), ) @@ -786,7 +945,7 @@ def activate(self) -> None: Client, "async_run", _wrap_async_run( - getattr(Client, "async_run"), + Client.async_run, span_name=f"{ASYNC_PREFIX}{REPLICATE_RUN_SPAN_NAME}", ), ) @@ -794,7 +953,7 @@ def activate(self) -> None: Client, "stream", _wrap_sync_run( - getattr(Client, "stream"), + Client.stream, span_name=REPLICATE_STREAM_SPAN_NAME, stream=True, ), @@ -803,7 +962,7 @@ def activate(self) -> None: Client, "async_stream", _wrap_async_run( - getattr(Client, "async_stream"), + Client.async_stream, span_name=f"{ASYNC_PREFIX}{REPLICATE_STREAM_SPAN_NAME}", stream=True, ), @@ -812,12 +971,12 @@ def activate(self) -> None: self._patch_attr( Predictions, "create", - _wrap_prediction_create(getattr(Predictions, "create")), + _wrap_prediction_create(Predictions.create), ) self._patch_attr( Predictions, "async_create", - _wrap_prediction_create(getattr(Predictions, "async_create"), is_async=True), + _wrap_prediction_create(Predictions.async_create, is_async=True), ) for method_name in ("list", "get", "cancel"): self._patch_attr( @@ -842,22 +1001,22 @@ def activate(self) -> None: self._patch_attr( Prediction, "wait", - _wrap_prediction_wait(getattr(Prediction, "wait")), + _wrap_prediction_wait(Prediction.wait), ) self._patch_attr( Prediction, "async_wait", - _wrap_prediction_wait(getattr(Prediction, "async_wait"), is_async=True), + _wrap_prediction_wait(Prediction.async_wait, is_async=True), ) self._patch_attr( Prediction, "stream", - _wrap_prediction_stream(getattr(Prediction, "stream")), + _wrap_prediction_stream(Prediction.stream), ) self._patch_attr( Prediction, "async_stream", - _wrap_prediction_stream(getattr(Prediction, "async_stream"), is_async=True), + _wrap_prediction_stream(Prediction.async_stream, is_async=True), ) self._patch_attr( @@ -882,15 +1041,51 @@ def activate(self) -> None: ) self._is_instrumented = True + + def activate(self) -> None: + """Monkey-patch the Replicate SDK with shared transactional ownership.""" + global _ENABLED, _REFCOUNT + + if self._is_instrumented: + return + with _LOCK: + if self._is_instrumented: + return + if _REFCOUNT: + _REFCOUNT += 1 + self._is_instrumented = True + return + try: + self._activate_once() + except Exception: + for patch in reversed(_PATCHES): + if getattr(patch.owner, patch.name, None) is patch.replacement: + setattr(patch.owner, patch.name, patch.original) + _PATCHES.clear() + raise + if not self._is_instrumented: + return + _ENABLED = True + _REFCOUNT = 1 logger.info("Replicate instrumentation activated") def deactivate(self) -> None: """Restore patched Replicate SDK methods.""" - for (owner, attr_name), original in reversed(self._originals.items()): - try: - setattr(owner, attr_name, original) - except Exception: - logger.debug("Failed to restore Replicate SDK attr %s", attr_name) - self._originals.clear() - self._is_instrumented = False + global _ENABLED, _REFCOUNT + + with _LOCK: + if not self._is_instrumented: + return + self._is_instrumented = False + _REFCOUNT = max(0, _REFCOUNT - 1) + if _REFCOUNT: + return + _ENABLED = False + for patch in reversed(_PATCHES): + try: + if getattr(patch.owner, patch.name, None) is patch.replacement: + setattr(patch.owner, patch.name, patch.original) + except Exception: # noqa: BLE001 - foreign-safe best-effort restore. + logger.debug("Failed to restore Replicate SDK attr %s", patch.name) + _PATCHES.clear() logger.info("Replicate instrumentation deactivated") diff --git a/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_serialization.py b/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_serialization.py new file mode 100644 index 00000000..7f039966 --- /dev/null +++ b/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_serialization.py @@ -0,0 +1,279 @@ +"""Bounded, privacy-safe serialization for Replicate SDK values.""" + +from __future__ import annotations + +import json +import math +import re +from collections.abc import Mapping, Sequence +from datetime import date, datetime +from enum import Enum +from itertools import islice +from numbers import Integral, Real +from typing import Any +from urllib.parse import urlsplit, urlunsplit + +MAX_ATTRIBUTE_BYTES = 16_000 +MAX_DEPTH = 8 +MAX_ITEMS = 50 +MAX_STRING_BYTES = 4_000 +REDACTED = "[REDACTED]" + +_SENSITIVE_SUFFIXES = ( + "apikey", + "authorization", + "credential", + "password", + "secret", + "sessiontoken", + "token", +) +_SECRET_ASSIGNMENT = re.compile( + r"(?i)([\"']?(?:api[_-]?key|authorization|password|secret|session[_-]?token|token)[\"']?)" + r"(\s*[:=]\s*)([\"']?)([^\s,;}\"']+)([\"']?)" +) +_BEARER = re.compile(r"(?i)\bbearer\s+[a-z0-9._~+/=-]+") +_ADDRESS = re.compile(r"(?i)\b0x[0-9a-f]{6,}\b") + + +def _truncate_utf8(value: str, limit: int = MAX_STRING_BYTES) -> str: + encoded = value.encode("utf-8") + if len(encoded) <= limit: + return value + suffix = "...[truncated]" + budget = max(0, limit - len(suffix.encode("utf-8"))) + return encoded[:budget].decode("utf-8", errors="ignore") + suffix + + +def sanitize_url(value: str) -> str: + try: + parsed = urlsplit(value) + except ValueError: + return "" + if not parsed.scheme or not parsed.netloc: + return value + hostname = parsed.hostname + if not hostname: + return "" + netloc = ( + f"[{hostname}]" + if ":" in hostname and not hostname.startswith("[") + else hostname + ) + try: + port = parsed.port + except ValueError: + return "" + if port is not None: + netloc = f"{netloc}:{port}" + return urlunsplit((parsed.scheme, netloc, parsed.path, "", "")) + + +def safe_text(value: Any, *, default: str = "") -> str: + if isinstance(value, str): + if "://" in value: + value = sanitize_url(value) + value = _BEARER.sub(REDACTED, value) + value = _SECRET_ASSIGNMENT.sub( + lambda match: ( + f"{match.group(1)}{match.group(2)}{match.group(3)}" + f"{REDACTED}{match.group(5)}" + ), + value, + ) + return _truncate_utf8(_ADDRESS.sub("0x", value)) + if value is None: + return default + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, Integral): + return str(int(value)) + if isinstance(value, Real) and math.isfinite(float(value)): + return str(float(value)) + return f"<{type(value).__name__}>" + + +def exception_message(exc: BaseException) -> str: + try: + arguments = exc.args + except Exception: # noqa: BLE001 + arguments = () + for argument in arguments: + if isinstance(argument, str | bool | int | float): + return safe_text(argument) + return type(exc).__name__ + + +def exception_status(exc: BaseException, *, default: int = 500) -> int: + for candidate in (exc, _safe_attr(exc, "response")): + for name in ("status_code", "status"): + value = _safe_attr(candidate, name) + if isinstance(value, int) and 400 <= value <= 599: + return value + return default + + +def _safe_attr(value: Any, name: str) -> Any: + try: + return getattr(value, name, None) + except Exception: # noqa: BLE001 + return None + + +def sensitive_key(value: Any) -> bool: + if not isinstance(value, str): + return False + normalized = re.sub(r"[^a-z0-9]", "", value.lower()) + return any(normalized.endswith(suffix) for suffix in _SENSITIVE_SUFFIXES) + + +def _key(value: Any) -> str: + if isinstance(value, Enum): + value = value.value + return safe_text(value)[:256] + + +def json_value(value: Any, *, depth: int = 0, seen: set[int] | None = None) -> Any: + if value is None or isinstance(value, bool): + return value + if isinstance(value, str): + return safe_text(value) + if isinstance(value, Integral): + return int(value) + if isinstance(value, Real): + number = float(value) + return number if math.isfinite(number) else None + if isinstance(value, datetime | date): + return value.isoformat() + if isinstance(value, bytes | bytearray | memoryview): + return {"length": len(value), "type": type(value).__name__} + if isinstance(value, Enum): + return json_value(value.value, depth=depth + 1, seen=seen) + if depth >= MAX_DEPTH: + return {"truncated": "max_depth", "type": type(value).__name__} + + active = seen if seen is not None else set() + identity = id(value) + if identity in active: + return "" + active.add(identity) + try: + if isinstance(value, Mapping): + result: dict[str, Any] = {} + items = list(islice(value.items(), MAX_ITEMS + 1)) + for key, item in items[:MAX_ITEMS]: + key_text = _key(key) + result[key_text] = ( + REDACTED + if sensitive_key(key) + else json_value(item, depth=depth + 1, seen=active) + ) + if len(items) > MAX_ITEMS: + result["__truncated_items__"] = True + return result + if isinstance(value, Sequence) and not isinstance( + value, str | bytes | bytearray + ): + items = list(islice(iter(value), MAX_ITEMS + 1)) + converted = [ + json_value(item, depth=depth + 1, seen=active) + for item in items[:MAX_ITEMS] + ] + if len(items) > MAX_ITEMS: + return {"items": converted, "truncated": True} + return converted + for method_name in ("model_dump", "to_dict", "dict"): + method = _safe_attr(value, method_name) + if not callable(method): + continue + try: + converted = method() + except Exception: # noqa: BLE001,S112 + continue + if isinstance(converted, Mapping): + return json_value(converted, depth=depth + 1, seen=active) + return {"type": type(value).__name__} + except Exception: # noqa: BLE001 + return {"type": type(value).__name__, "unserializable": True} + finally: + active.discard(identity) + + +def json_string(value: Any) -> str: + encoded = json.dumps( + json_value(value), + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + size = len(encoded.encode("utf-8")) + if size <= MAX_ATTRIBUTE_BYTES: + return encoded + low, high, result = 0, len(encoded), "" + while low <= high: + midpoint = (low + high) // 2 + candidate = json.dumps( + {"original_bytes": size, "preview": encoded[:midpoint], "truncated": True}, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + if len(candidate.encode("utf-8")) <= MAX_ATTRIBUTE_BYTES: + result = candidate + low = midpoint + 1 + else: + high = midpoint - 1 + return result + + +def prediction_summary(value: Any) -> Any: + """Return allow-listed stable prediction/page fields.""" + mapping = value if isinstance(value, Mapping) else None + if mapping is None: + for method_name in ("model_dump", "dict"): + method = _safe_attr(value, method_name) + if not callable(method): + continue + try: + converted = method() + except Exception: # noqa: BLE001,S112 + continue + if isinstance(converted, Mapping): + mapping = converted + break + if mapping is None: + results = _safe_attr(value, "results") + if results is not None: + return {"results": prediction_summary(results)} + return json_value(value) + + allowed = ( + "id", + "model", + "version", + "status", + "input", + "output", + "error", + "logs", + "metrics", + "created_at", + "started_at", + "completed_at", + "urls", + ) + summary = { + key: mapping[key] + for key in allowed + if key in mapping and mapping[key] is not None + } + if "results" in mapping: + summary["results"] = [ + prediction_summary(item) + for item in list(islice(iter(mapping["results"]), MAX_ITEMS)) + ] + for key in ("next", "previous"): + if key in mapping: + summary[key] = mapping[key] + return json_value(summary) diff --git a/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_translator.py b/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_translator.py index bec9dee1..b2a220fd 100644 --- a/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_translator.py +++ b/python-sdks/instrumentations/respan-instrumentation-replicate/src/respan_instrumentation_replicate/_translator.py @@ -2,11 +2,23 @@ from __future__ import annotations -import json from collections.abc import Mapping, Sequence +from itertools import islice from typing import Any from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes +from respan_sdk.constants.llm_logging import ( + LOG_TYPE_TASK, + LOG_TYPE_TEXT, + LogMethodChoices, +) +from respan_sdk.constants.span_attributes import ( + RESPAN_LOG_METHOD, + RESPAN_LOG_TYPE, + RESPAN_METADATA, + RESPAN_SPAN_ATTRIBUTES_MAP, + RESPAN_TRACE_GROUP_ID, +) from respan_instrumentation_replicate._constants import ( ASSISTANT_ROLE, @@ -19,8 +31,8 @@ METRICS_KEY, MODEL_KEY, OUTPUT_KEY, - PROMPT_KEY, PREDICTION_RESPAN_MODEL_ATTR, + PROMPT_KEY, REF_KEY, REPLICATE_SYSTEM_NAME, RESPAN_PARAMS_KEY, @@ -29,35 +41,22 @@ USER_ROLE, VERSION_KEY, ) -from respan_sdk.constants.llm_logging import ( - LOG_TYPE_TASK, - LOG_TYPE_TEXT, - LogMethodChoices, -) -from respan_sdk.constants.span_attributes import ( - RESPAN_LOG_METHOD, - RESPAN_LOG_TYPE, - RESPAN_METADATA, - RESPAN_SPAN_ATTRIBUTES_MAP, - RESPAN_TRACE_GROUP_ID, +from respan_instrumentation_replicate._serialization import ( + exception_message, + json_string, + prediction_summary, + safe_text, + sensitive_key, ) -from respan_sdk.utils.serialization import serialize_value def safe_json(value: Any) -> str: """Serialize arbitrary Replicate values into an OTEL-safe JSON string.""" - try: - return json.dumps( - serialize_value(value=value), default=str, separators=(",", ":") - ) - except Exception: - return json.dumps(str(value), separators=(",", ":")) + return json_string(value) def _truncate_text(value: str) -> str: - if len(value) <= MAX_TEXT_LENGTH: - return value - return f"{value[:MAX_TEXT_LENGTH]}..." + return safe_text(value)[:MAX_TEXT_LENGTH] def _to_mapping(value: Any) -> Mapping[str, Any] | None: @@ -65,26 +64,35 @@ def _to_mapping(value: Any) -> Mapping[str, Any] | None: return value for method_name in ("model_dump", "dict"): - method = getattr(value, method_name, None) + try: + method = getattr(value, method_name, None) + except Exception: # noqa: BLE001 - vendor objects may expose hostile properties. + method = None if callable(method): try: converted = method() - except Exception: + except Exception: # noqa: BLE001,S112 - vendor conversion is best effort. continue if isinstance(converted, Mapping): return converted - value_dict = getattr(value, "__dict__", None) + try: + value_dict = getattr(value, "__dict__", None) + except Exception: # noqa: BLE001 - fall back to a stable type summary. + value_dict = None if isinstance(value_dict, Mapping): return value_dict return None def _file_output_text(value: Any) -> str | None: - if value.__class__.__name__ != "FileOutput": + if type(value).__name__ != "FileOutput": return None - url = getattr(value, "url", None) - return str(url) if url else str(value) + try: + url = getattr(value, "url", None) + except Exception: # noqa: BLE001 + url = None + return safe_text(url) if url else f"<{type(value).__name__}>" def output_to_text(value: Any) -> str: @@ -102,14 +110,16 @@ def output_to_text(value: Any) -> str: if file_output is not None: return _truncate_text(file_output) - if isinstance(value, Sequence) and not isinstance( - value, (str, bytes, bytearray) - ): - if all(isinstance(item, str) for item in value): - return _truncate_text("".join(value)) - text_parts = [output_to_text(item) for item in value] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + items = list(islice(iter(value), 51)) + selected = items[:50] + if all(isinstance(item, str) for item in selected): + suffix = "...[truncated items]" if len(items) > 50 else "" + return _truncate_text("".join(selected) + suffix) + text_parts = [output_to_text(item) for item in selected] if any(text_parts): - return _truncate_text("".join(text_parts)) + suffix = "...[truncated items]" if len(items) > 50 else "" + return _truncate_text("".join(text_parts) + suffix) mapping = _to_mapping(value) if mapping is not None: @@ -118,7 +128,7 @@ def output_to_text(value: Any) -> str: return output_to_text(mapping[key]) return _truncate_text(safe_json(mapping)) - return _truncate_text(str(value)) + return f"<{type(value).__name__}>" def _prediction_mapping(prediction: Any) -> Mapping[str, Any]: @@ -136,14 +146,17 @@ def model_from_ref_or_prediction( if respan_params is not None: model_override = respan_params.get(RESPAN_PARAMS_MODEL_KEY) if model_override: - return str(model_override) + return safe_text(model_override) if prediction is not None: - prediction_model_override = getattr( - prediction, PREDICTION_RESPAN_MODEL_ATTR, None - ) + try: + prediction_model_override = getattr( + prediction, PREDICTION_RESPAN_MODEL_ATTR, None + ) + except Exception: # noqa: BLE001 - private attribute is best-effort. + prediction_model_override = None if prediction_model_override: - return str(prediction_model_override) + return safe_text(prediction_model_override) for value in ( kwargs.get(MODEL_KEY), @@ -152,21 +165,21 @@ def model_from_ref_or_prediction( ref, ): if value: - return str(value) + return safe_text(value) prediction_map = _prediction_mapping(prediction) model = prediction_map.get(MODEL_KEY) version = prediction_map.get(VERSION_KEY) if model and version: - model_text = str(model) - version_text = str(version) + model_text = safe_text(model) + version_text = safe_text(version) if version_text.startswith(f"{model_text}:"): return version_text return f"{model_text}:{version_text}" if model: - return str(model) + return safe_text(model) if version: - return str(version) + return safe_text(version) return None @@ -180,12 +193,14 @@ def _prompt_content(input_value: Any) -> str: return output_to_text(input_value) -def _base_llm_attrs(*, span_name: str, model: str | None, stream: bool) -> dict[str, Any]: +def _base_llm_attrs( + *, span_name: str, model: str | None, stream: bool +) -> dict[str, Any]: attrs: dict[str, Any] = { RESPAN_LOG_METHOD: LogMethodChoices.TRACING_INTEGRATION.value, RESPAN_LOG_TYPE: LOG_TYPE_TEXT, SpanAttributes.LLM_SYSTEM: REPLICATE_SYSTEM_NAME, - SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value, + SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value, SpanAttributes.TRACELOOP_ENTITY_NAME: span_name, SpanAttributes.TRACELOOP_ENTITY_PATH: span_name, } @@ -213,7 +228,7 @@ def _apply_respan_params(attributes: dict[str, Any], params: Any) -> str | None: span_name = params_mapping.get("span_name") workflow_name = params_mapping.get("workflow_name") if workflow_name and "trace_group_identifier" not in params_mapping: - attributes.setdefault(RESPAN_TRACE_GROUP_ID, str(workflow_name)) + attributes.setdefault(RESPAN_TRACE_GROUP_ID, safe_text(workflow_name)) for key, value in params_mapping.items(): if key in { @@ -229,15 +244,26 @@ def _apply_respan_params(attributes: dict[str, Any], params: Any) -> str | None: if attr_key is None: continue if attr_key == RESPAN_METADATA and isinstance(value, Mapping): + attributes[RESPAN_METADATA] = safe_json(value) for metadata_key, metadata_value in value.items(): - attributes[f"{RESPAN_METADATA}.{metadata_key}"] = ( - metadata_value - if isinstance(metadata_value, str) - else str(metadata_value) - ) + attribute_key = f"{RESPAN_METADATA}.{safe_text(metadata_key)[:128]}" + if sensitive_key(metadata_key): + attributes[attribute_key] = "[REDACTED]" + elif isinstance(metadata_value, str): + attributes[attribute_key] = safe_text(metadata_value) + elif isinstance(metadata_value, bool | int | float): + attributes[attribute_key] = metadata_value + else: + attributes[attribute_key] = safe_json(metadata_value) else: - attributes[attr_key] = value - return str(span_name) if span_name else None + attributes[attr_key] = ( + safe_text(value) + if isinstance(value, str | bytes | bytearray) + else value + if isinstance(value, bool | int | float) + else safe_json(value) + ) + return safe_text(span_name) if span_name else None def _set_request_attrs( @@ -270,6 +296,9 @@ def _set_prediction_metadata(attrs: dict[str, Any], prediction: Any) -> None: metrics = prediction_map.get(METRICS_KEY) if metrics: attrs[f"{RESPAN_METADATA}.replicate_metrics"] = safe_json(metrics) + attrs[f"{RESPAN_METADATA}.replicate_prediction"] = safe_json( + prediction_summary(prediction) + ) def _set_output_attrs( @@ -280,7 +309,7 @@ def _set_output_attrs( error: Exception | None, ) -> None: if error is not None: - completion_text = str(error) + completion_text = exception_message(error) elif output is not None: completion_text = output_to_text(output) else: @@ -288,7 +317,7 @@ def _set_output_attrs( completion_text = output_to_text(prediction_map.get(OUTPUT_KEY)) if completion_text: - attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = completion_text + attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = safe_json(completion_text) completion_prefix = f"{SpanAttributes.LLM_COMPLETIONS}.0" attrs[f"{completion_prefix}.role"] = ASSISTANT_ROLE attrs[f"{completion_prefix}.content"] = completion_text @@ -334,7 +363,14 @@ def build_operation_span_data( if input_value is not None: attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] = safe_json(input_value) if output is not None: - attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = safe_json(output) + attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = safe_json( + prediction_summary(output) + ) if error is not None: - attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = str(error) + attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] = safe_json( + { + "error": type(error).__name__, + "message": exception_message(error), + } + ) return span_name, attrs diff --git a/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_instrumentation.py index 5d9b1225..3c7d43ac 100644 --- a/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_instrumentation.py @@ -4,9 +4,7 @@ import pytest from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes - -from respan_instrumentation_replicate import ReplicateInstrumentor -from respan_instrumentation_replicate import _instrumentation +from respan_instrumentation_replicate import ReplicateInstrumentor, _instrumentation from respan_instrumentation_replicate._constants import ( OFF_CONTRACT_ALIASES, RESPAN_PARAMS_KEY, @@ -27,8 +25,11 @@ @pytest.fixture(autouse=True) -def reset_tracer(): +def reset_tracer(monkeypatch): RespanTracer.reset_instance() + monkeypatch.setattr(_instrumentation, "_REFCOUNT", 0) + monkeypatch.setattr(_instrumentation, "_ENABLED", False) + monkeypatch.setattr(_instrumentation, "_PATCHES", []) yield RespanTracer.reset_instance() @@ -51,7 +52,7 @@ def test_build_model_call_span_data_uses_canonical_attrs_only(): assert span_name == "replicate.run" assert attrs[RESPAN_LOG_TYPE] == LOG_TYPE_TEXT assert attrs[SpanAttributes.LLM_SYSTEM] == "replicate" - assert attrs[SpanAttributes.LLM_REQUEST_TYPE] == LLMRequestTypeValues.COMPLETION.value + assert attrs[SpanAttributes.LLM_REQUEST_TYPE] == LLMRequestTypeValues.CHAT.value assert attrs[SpanAttributes.LLM_REQUEST_MODEL] == "meta/meta-llama-3-8b-instruct" assert attrs[f"{SpanAttributes.LLM_PROMPTS}.0.role"] == "user" assert attrs[f"{SpanAttributes.LLM_PROMPTS}.0.content"] == "Say hi" @@ -96,7 +97,9 @@ def test_model_from_prediction_does_not_duplicate_prefixed_version(): model="owner/model", version="owner/model:version-id", ) - assert model_from_ref_or_prediction(prediction=prediction) == "owner/model:version-id" + assert ( + model_from_ref_or_prediction(prediction=prediction) == "owner/model:version-id" + ) def _install_fake_replicate_modules(monkeypatch): @@ -230,8 +233,14 @@ def test_instrumentor_patches_run_and_stream(monkeypatch): assert len(fake.emitted_spans) == 3 assert fake.emitted_spans[0].name == "replicate.run" - assert fake.emitted_spans[0].attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] == "hello world" - assert fake.emitted_spans[1].attributes[SpanAttributes.LLM_REQUEST_MODEL] == "gpt-4o-mini" + assert ( + fake.emitted_spans[0].attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] + == "hello world" + ) + assert ( + fake.emitted_spans[1].attributes[SpanAttributes.LLM_REQUEST_MODEL] + == "gpt-4o-mini" + ) assert fake.emitted_spans[2].name == "replicate.stream" assert fake.emitted_spans[2].attributes[SpanAttributes.LLM_IS_STREAMING] is True @@ -266,8 +275,14 @@ async def run_calls(): "replicate.predictions.create", "replicate.prediction.wait", ] - assert fake.emitted_spans[1].attributes[SpanAttributes.LLM_REQUEST_MODEL] == "gpt-4o-mini" - assert fake.emitted_spans[2].attributes[SpanAttributes.LLM_REQUEST_MODEL] == "gpt-4o-mini" + assert ( + fake.emitted_spans[1].attributes[SpanAttributes.LLM_REQUEST_MODEL] + == "gpt-4o-mini" + ) + wait_attrs = fake.emitted_spans[2].attributes + assert wait_attrs[RESPAN_LOG_TYPE] == "task" + assert SpanAttributes.LLM_REQUEST_MODEL not in wait_attrs + assert SpanAttributes.LLM_SYSTEM not in wait_attrs def test_activate_skips_when_respan_tracing_is_disabled(monkeypatch, caplog): @@ -279,7 +294,10 @@ def test_activate_skips_when_respan_tracing_is_disabled(monkeypatch, caplog): instrumentor.activate() assert instrumentor._is_instrumented is False - assert "Replicate instrumentation skipped because Respan tracing is disabled" in caplog.text + assert ( + "Replicate instrumentation skipped because Respan tracing is disabled" + in caplog.text + ) assert fake.client_class().run("owner/model", input={"prompt": "hi"}) == [ "hello", " world", diff --git a/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_real_replicate_runtime.py b/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_real_replicate_runtime.py new file mode 100644 index 00000000..fb2df437 --- /dev/null +++ b/python-sdks/instrumentations/respan-instrumentation-replicate/tests/test_real_replicate_runtime.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import json +from collections import Counter + +import httpx +import pytest +import replicate +from opentelemetry.semconv_ai import SpanAttributes +from replicate.exceptions import ReplicateError +from respan_instrumentation_replicate import ReplicateInstrumentor +from respan_instrumentation_replicate import _instrumentation as instrumentation +from respan_instrumentation_replicate._serialization import json_string +from respan_sdk.constants.span_attributes import RESPAN_LOG_TYPE +from respan_tracing.core.tracer import RespanTracer + + +def _prediction(identifier: str = "pred-1", *, prompt: str = "hello") -> dict: + return { + "id": identifier, + "model": "owner/model", + "version": "owner/model:version-1", + "status": "succeeded", + "input": {"prompt": prompt}, + "output": ["bounded ", "answer"], + "error": None, + "logs": "completed", + "metrics": {"predict_time": 0.01}, + "urls": { + "get": f"https://mock.replicate.local/v1/predictions/{identifier}", + "cancel": f"https://mock.replicate.local/v1/predictions/{identifier}/cancel", + }, + "created_at": "2026-08-18T00:00:00Z", + "started_at": "2026-08-18T00:00:00Z", + "completed_at": "2026-08-18T00:00:01Z", + } + + +@pytest.fixture(autouse=True) +def clean_runtime(monkeypatch): + RespanTracer.reset_instance() + monkeypatch.setattr(instrumentation, "_REFCOUNT", 0) + monkeypatch.setattr(instrumentation, "_PATCHES", []) + monkeypatch.setattr(instrumentation, "_ENABLED", False) + yield + RespanTracer.reset_instance() + + +def test_real_current_sdk_exports_run_wait_and_management_without_duplicates( + monkeypatch, +) -> None: + emitted = [] + monkeypatch.setattr( + instrumentation, + "inject_span", + lambda span: emitted.append(span) or True, + ) + + def handler(request: httpx.Request) -> httpx.Response: + path = request.url.path + if request.method == "POST" and path.endswith("/predictions"): + body = json.loads(request.content) + return httpx.Response( + 201, + json=_prediction(prompt=(body.get("input") or {}).get("prompt", "")), + ) + if request.method == "GET" and path == "/v1/predictions": + return httpx.Response( + 200, + json={"next": None, "previous": None, "results": [_prediction()]}, + ) + if request.method == "GET" and path == "/v1/predictions/pred-1": + return httpx.Response(200, json=_prediction()) + return httpx.Response( + 404, json={"detail": f"unhandled {request.method} {path}"} + ) + + adapter = ReplicateInstrumentor() + adapter.activate() + client = replicate.Client( + api_token="fixture-token", + base_url="https://mock.replicate.local", + transport=httpx.MockTransport(handler), + ) + try: + assert client.run("owner/model", input={"prompt": "hello"}) == [ + "bounded ", + "answer", + ] + prediction = client.predictions.create( + version="owner/model:version-1", + input={"prompt": "lifecycle"}, + ) + prediction.wait() + assert client.predictions.get("pred-1").id == "pred-1" + assert len(client.predictions.list().results) == 1 + finally: + adapter.deactivate() + + names = Counter(span.name for span in emitted) + assert names == Counter( + { + "replicate.run": 1, + "replicate.predictions.create": 1, + "replicate.prediction.wait": 1, + "replicate.predictions.get": 1, + "replicate.predictions.list": 1, + } + ) + assert len({span.context.span_id for span in emitted}) == 5 + run = next(span for span in emitted if span.name == "replicate.run") + wait = next(span for span in emitted if span.name == "replicate.prediction.wait") + listed = next(span for span in emitted if span.name == "replicate.predictions.list") + assert run.attributes[RESPAN_LOG_TYPE] == "text" + assert run.attributes[SpanAttributes.LLM_REQUEST_TYPE] == "chat" + assert wait.attributes[RESPAN_LOG_TYPE] == "task" + assert SpanAttributes.LLM_REQUEST_MODEL not in wait.attributes + assert listed.attributes[RESPAN_LOG_TYPE] == "task" + listed_output = listed.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT] + assert "__orig_class__" not in listed_output + assert "0x" not in listed_output + assert all("traceloop.span.kind" not in span.attributes for span in emitted) + assert all(span.instrumentation_scope.name == "replicate" for span in emitted) + + +def test_real_current_sdk_preserves_provider_status_and_safe_error(monkeypatch) -> None: + emitted = [] + monkeypatch.setattr( + instrumentation, + "inject_span", + lambda span: emitted.append(span) or True, + ) + + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 429, + json={"detail": 'api_key="plain-secret" provider limited'}, + ) + + adapter = ReplicateInstrumentor() + adapter.activate() + client = replicate.Client( + api_token="fixture-token", + base_url="https://mock.replicate.local", + transport=httpx.MockTransport(handler), + ) + try: + with pytest.raises(ReplicateError): + client.run("owner/model", input={"prompt": "fail"}) + finally: + adapter.deactivate() + + assert len(emitted) == 1 + span = emitted[0] + assert span.attributes["status_code"] == 429 + assert span.status.status_code.name == "ERROR" + assert "plain-secret" not in span.attributes["error.message"] + + +def test_replicate_serialization_is_valid_bounded_and_private() -> None: + class Hostile: + def __str__(self) -> str: + raise AssertionError("must not stringify") + + def __repr__(self) -> str: + raise AssertionError("must not repr") + + encoded = json_string( + { + "auth_token": "plain-secret", + "content": "😀" * 5_000, + "hostile": Hostile(), + "nonfinite": float("nan"), + } + ) + assert len(encoded.encode("utf-8")) <= 16_000 + assert "plain-secret" not in encoded + assert json.loads(encoded)["nonfinite"] is None diff --git a/python-sdks/instrumentations/respan-instrumentation-restate/README.md b/python-sdks/instrumentations/respan-instrumentation-restate/README.md index 59ec2849..7bbca04e 100644 --- a/python-sdks/instrumentations/respan-instrumentation-restate/README.md +++ b/python-sdks/instrumentations/respan-instrumentation-restate/README.md @@ -28,6 +28,7 @@ respan = Respan( greeter = restate.Service("Greeter") + @greeter.handler() async def greet(ctx: restate.Context, name: str) -> str: return f"Hello, {name}!" diff --git a/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_instrumentation.py index 09b2f3e4..ac50cf5b 100644 --- a/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_instrumentation.py @@ -3,24 +3,16 @@ from __future__ import annotations import importlib -import json +import importlib.metadata import logging import threading -from collections.abc import Mapping, Sequence from contextlib import asynccontextmanager +from dataclasses import dataclass from typing import Any from opentelemetry import trace -from opentelemetry.instrumentation.utils import unwrap -from opentelemetry.trace import Status, StatusCode from opentelemetry.semconv_ai import SpanAttributes -from wrapt import wrap_function_wrapper - -from respan_instrumentation_restate._constants import ( - RESTATE_CONTEXT_MANAGER_MARKER, - RESTATE_INSTRUMENTATION_NAME, - RESTATE_REGISTRATION_TARGETS, -) +from opentelemetry.trace import Status, StatusCode from respan_sdk.constants import ERROR_MESSAGE_ATTR from respan_sdk.constants.llm_logging import LogMethodChoices from respan_sdk.constants.span_attributes import ( @@ -30,18 +22,43 @@ RESPAN_THREADS_ID, RESPAN_TRACE_GROUP_ID, ) -from respan_sdk.utils.serialization import serialize_value from respan_tracing.core.tracer import RespanTracer +from respan_tracing.utils.span_factory import read_propagated_attributes +from wrapt import FunctionWrapper + +from respan_instrumentation_restate._constants import ( + RESTATE_CONTEXT_MANAGER_MARKER, + RESTATE_INSTRUMENTATION_NAME, + RESTATE_REGISTRATION_TARGETS, +) +from respan_instrumentation_restate._serialization import ( + exception_message, + exception_status, + json_string, + json_value, + safe_text, + sensitive_key, +) logger = logging.getLogger(__name__) _LOCK = threading.RLock() _ACTIVATION_COUNT = 0 -_PATCHED_TARGETS: list[tuple[str, str]] = [] _ENABLED = False _CAPTURE_CONTENT = True +@dataclass +class _Patch: + owner: Any + name: str + original: Any + replacement: Any + + +_PATCHED_TARGETS: list[_Patch] = [] + + def _is_respan_tracing_enabled() -> bool: tracer = getattr(RespanTracer, "_instance", None) if tracer is None: @@ -49,65 +66,49 @@ def _is_respan_tracing_enabled() -> bool: return bool(getattr(tracer, "is_enabled", True)) -def _jsonable(value: Any, *, depth: int = 0) -> Any: - if depth > 8: - return repr(value) - if value is None or isinstance(value, (str, int, float, bool)): - return value - if isinstance(value, bytes): - try: - return value.decode("utf-8") - except UnicodeDecodeError: - return {"type": "bytes", "length": len(value)} - if isinstance(value, Mapping): - return { - str(key): _jsonable(item, depth=depth + 1) for key, item in value.items() - } - if isinstance(value, Sequence): - return [_jsonable(item, depth=depth + 1) for item in value] +def _safe_attr(value: Any, name: str, default: Any = None) -> Any: try: - return serialize_value(value=value) - except Exception: - return repr(value) - - -def _json_string(value: Any) -> str: - return json.dumps(_jsonable(value), default=str, ensure_ascii=False) + return getattr(value, name, default) + except Exception: # noqa: BLE001 - Restate objects are untrusted input. + return default def _deserialize_input(context: Any) -> Any: - invocation = context.invocation - handler_io = context.handler.handler_io + invocation = _safe_attr(context, "invocation") + handler = _safe_attr(context, "handler") + handler_io = _safe_attr(handler, "handler_io") try: return handler_io.input_serde.deserialize(invocation.input_buffer) - except Exception: - return _jsonable(invocation.input_buffer) + except Exception: # noqa: BLE001 - invalid input falls back to a typed summary. + return json_value(_safe_attr(invocation, "input_buffer")) def _invocation_details(context: Any) -> tuple[dict[str, Any], dict[str, Any]]: - handler = context.handler - service_tag = handler.service_tag - invocation = context.invocation + handler = _safe_attr(context, "handler") + service_tag = _safe_attr(handler, "service_tag") + invocation = _safe_attr(context, "invocation") server_context = importlib.import_module("restate.server_context") replaying_var = getattr(server_context, "restate_context_is_replaying", None) is_replaying = bool(replaying_var.get()) if replaying_var is not None else False metadata = { - "service_kind": service_tag.kind, - "service_name": service_tag.name, - "handler_name": handler.name, - "handler_kind": handler.kind, - "invocation_id": invocation.invocation_id, + "service_kind": safe_text(_safe_attr(service_tag, "kind")), + "service_name": safe_text(_safe_attr(service_tag, "name")), + "handler_name": safe_text(_safe_attr(handler, "name")), + "handler_kind": safe_text(_safe_attr(handler, "kind")), + "invocation_id": safe_text(_safe_attr(invocation, "invocation_id")), "replaying": is_replaying, } for name in ("key", "scope", "limit_key", "idempotency_key"): - value = getattr(invocation, name, None) + value = _safe_attr(invocation, name) if value: - metadata[name] = value - if service_tag.metadata: - metadata["service_metadata"] = dict(service_tag.metadata) - if handler.metadata: - metadata["handler_metadata"] = dict(handler.metadata) + metadata[name] = safe_text(value) + service_metadata = _safe_attr(service_tag, "metadata") + if service_metadata: + metadata["service_metadata"] = json_value(service_metadata) + handler_metadata = _safe_attr(handler, "metadata") + if handler_metadata: + metadata["handler_metadata"] = json_value(handler_metadata) input_payload = dict(metadata) if _CAPTURE_CONTENT: @@ -116,8 +117,9 @@ def _invocation_details(context: Any) -> tuple[dict[str, Any], dict[str, Any]]: def _log_type(context: Any) -> str: - service_kind = str(context.handler.service_tag.kind or "") - handler_kind = str(context.handler.kind or "") + handler = _safe_attr(context, "handler") + service_kind = safe_text(_safe_attr(_safe_attr(handler, "service_tag"), "kind")) + handler_kind = safe_text(_safe_attr(handler, "kind")) return ( "workflow" if service_kind == "workflow" and handler_kind == "workflow" @@ -131,21 +133,48 @@ def _span_attributes( metadata: dict[str, Any], input_payload: dict[str, Any], ) -> dict[str, Any]: - handler = context.handler - invocation = context.invocation - entity_name = f"{handler.service_tag.name}.{handler.name}" + handler = _safe_attr(context, "handler") + invocation = _safe_attr(context, "invocation") + service_name = safe_text(_safe_attr(_safe_attr(handler, "service_tag"), "name")) + handler_name = safe_text(_safe_attr(handler, "name")) + entity_name = f"{service_name}.{handler_name}" + current = trace.get_current_span() + try: + has_parent = bool(current.get_span_context().is_valid) + except Exception: # noqa: BLE001 + has_parent = False attrs: dict[str, Any] = { RESPAN_LOG_METHOD: LogMethodChoices.TRACING_INTEGRATION.value, RESPAN_LOG_TYPE: _log_type(context), - RESPAN_TRACE_GROUP_ID: str(invocation.invocation_id), - RESPAN_METADATA: _json_string({"restate": metadata}), + RESPAN_TRACE_GROUP_ID: safe_text(_safe_attr(invocation, "invocation_id")), + RESPAN_METADATA: json_string({"restate": metadata}), SpanAttributes.TRACELOOP_ENTITY_NAME: entity_name, - SpanAttributes.TRACELOOP_ENTITY_PATH: entity_name, - SpanAttributes.TRACELOOP_ENTITY_INPUT: _json_string(input_payload), + SpanAttributes.TRACELOOP_ENTITY_PATH: entity_name if has_parent else "", } - key = getattr(invocation, "key", None) + if _CAPTURE_CONTENT: + attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] = json_string(input_payload) + key = _safe_attr(invocation, "key") if key: - attrs[RESPAN_THREADS_ID] = str(key) + attrs[RESPAN_THREADS_ID] = safe_text(key) + + propagated = read_propagated_attributes() + aggregate: dict[str, Any] = {"restate": metadata} + for attr_key, value in propagated.items(): + if attr_key == RESPAN_METADATA: + continue + if attr_key.startswith(f"{RESPAN_METADATA}."): + metadata_key = attr_key.removeprefix(f"{RESPAN_METADATA}.") + safe_value = ( + "[REDACTED]" if sensitive_key(metadata_key) else json_value(value) + ) + aggregate[metadata_key] = safe_value + if isinstance(safe_value, str | bool | int | float): + attrs[attr_key] = safe_value + else: + attrs[attr_key] = json_string(safe_value) + elif attr_key in {RESPAN_TRACE_GROUP_ID, RESPAN_THREADS_ID}: + attrs[attr_key] = safe_text(value) + attrs[RESPAN_METADATA] = json_string(aggregate) return attrs @@ -169,7 +198,11 @@ async def _invocation_context(): metadata=metadata, input_payload=input_payload, ) - tracer = trace.get_tracer(__name__) + try: + version = importlib.metadata.version("respan-instrumentation-restate") + except importlib.metadata.PackageNotFoundError: + version = None + tracer = trace.get_tracer(RESTATE_INSTRUMENTATION_NAME, version) span_name = ( f"restate.{context.handler.service_tag.kind}." f"{context.handler.service_tag.name}.{context.handler.name}" @@ -183,32 +216,37 @@ async def _invocation_context(): try: yield except BaseException as exc: - message = str(exc) or type(exc).__name__ - span.record_exception(exc) + message = exception_message(exc) span.set_status(Status(StatusCode.ERROR, message)) - span.set_attribute( - "status_code", - int(getattr(exc, "status_code", 500) or 500), - ) + span.set_attribute("status_code", exception_status(exc)) span.set_attribute(ERROR_MESSAGE_ATTR, message) - span.set_attribute( - SpanAttributes.TRACELOOP_ENTITY_OUTPUT, - _json_string( - { - "status": "error", - "error": type(exc).__name__, - "message": message if _CAPTURE_CONTENT else type(exc).__name__, - } - ), + if _CAPTURE_CONTENT: + span.set_attribute( + SpanAttributes.TRACELOOP_ENTITY_OUTPUT, + json_string( + { + "status": "error", + "error": type(exc).__name__, + "message": message, + } + ), + ) + span.add_event( + "exception", + { + "exception.type": f"{type(exc).__module__}.{type(exc).__name__}", + "exception.message": message, + }, ) raise else: span.set_status(Status(StatusCode.OK)) span.set_attribute("status_code", 200) - span.set_attribute( - SpanAttributes.TRACELOOP_ENTITY_OUTPUT, - _json_string({"status": "completed"}), - ) + if _CAPTURE_CONTENT: + span.set_attribute( + SpanAttributes.TRACELOOP_ENTITY_OUTPUT, + json_string({"status": "completed"}), + ) setattr(_invocation_context, RESTATE_CONTEXT_MANAGER_MARKER, True) @@ -234,6 +272,29 @@ def _registration_wrapper( return wrapped(*args, **kwargs) +def _install_patches() -> None: + for module_path, target in RESTATE_REGISTRATION_TARGETS: + module = importlib.import_module(module_path) + owner: Any = module + owner_path, name = target.rsplit(".", maxsplit=1) + for component in owner_path.split("."): + owner = getattr(owner, component) + original = getattr(owner, name) + replacement = FunctionWrapper(original, _registration_wrapper) + setattr(owner, name, replacement) + _PATCHED_TARGETS.append(_Patch(owner, name, original, replacement)) + + +def _remove_patches() -> None: + for patch in reversed(_PATCHED_TARGETS): + try: + if getattr(patch.owner, patch.name, None) is patch.replacement: + setattr(patch.owner, patch.name, patch.original) + except Exception: # noqa: BLE001 - foreign-safe best-effort restore. + logger.debug("Failed to restore Restate target %s", patch.name) + _PATCHED_TARGETS.clear() + + class RestateInstrumentor: """Inject canonical Respan spans into Restate handler invocations.""" @@ -259,30 +320,15 @@ def activate(self) -> None: if _ACTIVATION_COUNT == 0: _CAPTURE_CONTENT = self._capture_content try: - for module_path, target in RESTATE_REGISTRATION_TARGETS: - wrap_function_wrapper( - module_path, - target, - _registration_wrapper, - ) - _PATCHED_TARGETS.append((module_path, target)) + _install_patches() except Exception: - for module_path, target in reversed(_PATCHED_TARGETS): - try: - unwrap(module_path, target) - except Exception: - logger.debug( - "Failed to roll back %s.%s", - module_path, - target, - exc_info=True, - ) - _PATCHED_TARGETS.clear() + _remove_patches() raise _ENABLED = True elif _CAPTURE_CONTENT != self._capture_content: - logger.warning( - "Restate is already active; the first capture_content setting wins" + raise ValueError( + "all active RestateInstrumentor instances must use the same " + "capture_content setting" ) _ACTIVATION_COUNT += 1 self._is_instrumented = True @@ -299,14 +345,4 @@ def deactivate(self) -> None: if _ACTIVATION_COUNT: return _ENABLED = False - for module_path, target in reversed(_PATCHED_TARGETS): - try: - unwrap(module_path, target) - except Exception: - logger.debug( - "Failed to unwrap %s.%s", - module_path, - target, - exc_info=True, - ) - _PATCHED_TARGETS.clear() + _remove_patches() diff --git a/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_serialization.py b/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_serialization.py new file mode 100644 index 00000000..38270ea9 --- /dev/null +++ b/python-sdks/instrumentations/respan-instrumentation-restate/src/respan_instrumentation_restate/_serialization.py @@ -0,0 +1,214 @@ +"""Bounded, privacy-safe serialization for Restate invocation data.""" + +from __future__ import annotations + +import json +import math +import re +from collections.abc import Mapping, Sequence +from enum import Enum +from itertools import islice +from numbers import Integral, Real +from typing import Any +from urllib.parse import urlsplit, urlunsplit + +MAX_ATTRIBUTE_BYTES = 16_000 +MAX_DEPTH = 8 +MAX_ITEMS = 50 +MAX_STRING_BYTES = 4_000 +REDACTED = "[REDACTED]" + +_SENSITIVE_SUFFIXES = ( + "apikey", + "authorization", + "credential", + "password", + "secret", + "sessiontoken", + "token", +) +_SECRET_ASSIGNMENT = re.compile( + r"(?i)([\"']?(?:api[_-]?key|authorization|password|secret|session[_-]?token|token)[\"']?)" + r"(\s*[:=]\s*)([\"']?)([^\s,;}\"']+)([\"']?)" +) +_BEARER = re.compile(r"(?i)\bbearer\s+[a-z0-9._~+/=-]+") +_ADDRESS = re.compile(r"(?i)\b0x[0-9a-f]{6,}\b") + + +def _truncate_utf8(value: str, limit: int = MAX_STRING_BYTES) -> str: + encoded = value.encode("utf-8") + if len(encoded) <= limit: + return value + suffix = "...[truncated]" + budget = max(0, limit - len(suffix.encode("utf-8"))) + return encoded[:budget].decode("utf-8", errors="ignore") + suffix + + +def sanitize_url(value: str) -> str: + try: + parsed = urlsplit(value) + except ValueError: + return "" + if not parsed.scheme or not parsed.netloc: + return value + hostname = parsed.hostname + if not hostname: + return "" + netloc = ( + f"[{hostname}]" + if ":" in hostname and not hostname.startswith("[") + else hostname + ) + try: + port = parsed.port + except ValueError: + return "" + if port is not None: + netloc = f"{netloc}:{port}" + return urlunsplit((parsed.scheme, netloc, parsed.path, "", "")) + + +def safe_text(value: Any, *, default: str = "") -> str: + if isinstance(value, str): + if "://" in value: + value = sanitize_url(value) + value = _BEARER.sub(REDACTED, value) + value = _SECRET_ASSIGNMENT.sub( + lambda match: ( + f"{match.group(1)}{match.group(2)}{match.group(3)}" + f"{REDACTED}{match.group(5)}" + ), + value, + ) + return _truncate_utf8(_ADDRESS.sub("0x", value)) + if value is None: + return default + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, Integral): + return str(int(value)) + if isinstance(value, Real) and math.isfinite(float(value)): + return str(float(value)) + return f"<{type(value).__name__}>" + + +def exception_message(exc: BaseException) -> str: + try: + arguments = exc.args + except Exception: # noqa: BLE001 + arguments = () + for argument in arguments: + if isinstance(argument, str | bool | int | float): + return safe_text(argument) + return type(exc).__name__ + + +def exception_status(exc: BaseException, *, default: int = 500) -> int: + try: + response = getattr(exc, "response", None) + except Exception: # noqa: BLE001 + response = None + for candidate in (exc, response): + for name in ("status_code", "status"): + try: + value = getattr(candidate, name, None) + except Exception: # noqa: BLE001 + value = None + if isinstance(value, int) and 400 <= value <= 599: + return value + return default + + +def sensitive_key(value: Any) -> bool: + if not isinstance(value, str): + return False + normalized = re.sub(r"[^a-z0-9]", "", value.lower()) + return any(normalized.endswith(suffix) for suffix in _SENSITIVE_SUFFIXES) + + +def _key(value: Any) -> str: + if isinstance(value, Enum): + value = value.value + return safe_text(value)[:256] + + +def json_value(value: Any, *, depth: int = 0, seen: set[int] | None = None) -> Any: + if value is None or isinstance(value, bool): + return value + if isinstance(value, str): + return safe_text(value) + if isinstance(value, Integral): + return int(value) + if isinstance(value, Real): + number = float(value) + return number if math.isfinite(number) else None + if isinstance(value, bytes | bytearray | memoryview): + return {"length": len(value), "type": type(value).__name__} + if isinstance(value, Enum): + return json_value(value.value, depth=depth + 1, seen=seen) + if depth >= MAX_DEPTH: + return {"truncated": "max_depth", "type": type(value).__name__} + + active = seen if seen is not None else set() + identity = id(value) + if identity in active: + return "" + active.add(identity) + try: + if isinstance(value, Mapping): + result: dict[str, Any] = {} + items = list(islice(value.items(), MAX_ITEMS + 1)) + for key, item in items[:MAX_ITEMS]: + key_text = _key(key) + result[key_text] = ( + REDACTED + if sensitive_key(key) + else json_value(item, depth=depth + 1, seen=active) + ) + if len(items) > MAX_ITEMS: + result["__truncated_items__"] = True + return result + if isinstance(value, Sequence) and not isinstance( + value, str | bytes | bytearray + ): + items = list(islice(iter(value), MAX_ITEMS + 1)) + converted = [ + json_value(item, depth=depth + 1, seen=active) + for item in items[:MAX_ITEMS] + ] + if len(items) > MAX_ITEMS: + return {"items": converted, "truncated": True} + return converted + return {"type": type(value).__name__} + except Exception: # noqa: BLE001 + return {"type": type(value).__name__, "unserializable": True} + finally: + active.discard(identity) + + +def json_string(value: Any) -> str: + encoded = json.dumps( + json_value(value), + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + size = len(encoded.encode("utf-8")) + if size <= MAX_ATTRIBUTE_BYTES: + return encoded + low, high, result = 0, len(encoded), "" + while low <= high: + midpoint = (low + high) // 2 + candidate = json.dumps( + {"original_bytes": size, "preview": encoded[:midpoint], "truncated": True}, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + if len(candidate.encode("utf-8")) <= MAX_ATTRIBUTE_BYTES: + result = candidate + low = midpoint + 1 + else: + high = midpoint - 1 + return result diff --git a/python-sdks/instrumentations/respan-instrumentation-restate/tests/test_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-restate/tests/test_instrumentation.py index c6ea54ac..e5a4d9cb 100644 --- a/python-sdks/instrumentations/respan-instrumentation-restate/tests/test_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-restate/tests/test_instrumentation.py @@ -1,15 +1,23 @@ import asyncio +import json from contextlib import contextmanager from types import SimpleNamespace -from respan_instrumentation_restate import RestateInstrumentor +import pytest import respan_instrumentation_restate._instrumentation as instrumentation +import restate +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.semconv_ai import SpanAttributes +from opentelemetry.trace import StatusCode +from respan_instrumentation_restate import RestateInstrumentor +from respan_instrumentation_restate._serialization import json_string from respan_sdk.constants.span_attributes import ( RESPAN_LOG_TYPE, RESPAN_THREADS_ID, RESPAN_TRACE_GROUP_ID, ) -from opentelemetry.semconv_ai import SpanAttributes class FakeSpan: @@ -17,6 +25,7 @@ def __init__(self, attributes: dict): self.attributes = dict(attributes) self.status = None self.exceptions = [] + self.events = [] def set_attribute(self, key, value): self.attributes[key] = value @@ -27,6 +36,9 @@ def set_status(self, status): def record_exception(self, exception): self.exceptions.append(exception) + def add_event(self, name, attributes): + self.events.append((name, attributes)) + class FakeTracer: def __init__(self): @@ -83,7 +95,9 @@ def import_module(name: str): return real_import(name) monkeypatch.setattr(instrumentation.importlib, "import_module", import_module) - monkeypatch.setattr(instrumentation.trace, "get_tracer", lambda name: fake_tracer) + monkeypatch.setattr( + instrumentation.trace, "get_tracer", lambda *args, **kwargs: fake_tracer + ) monkeypatch.setattr(instrumentation, "_ENABLED", True) monkeypatch.setattr(instrumentation, "_CAPTURE_CONTENT", True) @@ -96,8 +110,8 @@ async def run(): assert attrs[RESPAN_LOG_TYPE] == "workflow" assert attrs[RESPAN_TRACE_GROUP_ID] == "inv-123" assert attrs[RESPAN_THREADS_ID] == "order-42" - assert '"input": {"name": "Ada"}' in attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] - assert '"replaying": true' in attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] + assert '"input":{"name":"Ada"}' in attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] + assert '"replaying":true' in attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] assert attrs["status_code"] == 200 @@ -115,7 +129,9 @@ def import_module(name: str): return real_import(name) monkeypatch.setattr(instrumentation.importlib, "import_module", import_module) - monkeypatch.setattr(instrumentation.trace, "get_tracer", lambda name: fake_tracer) + monkeypatch.setattr( + instrumentation.trace, "get_tracer", lambda *args, **kwargs: fake_tracer + ) monkeypatch.setattr(instrumentation, "_ENABLED", True) async def run(): @@ -129,7 +145,7 @@ async def run(): attrs = fake_tracer.span.attributes assert attrs["status_code"] == 500 assert attrs["error.message"] == "deterministic Restate failure" - assert len(fake_tracer.span.exceptions) == 1 + assert fake_tracer.span.events[0][0] == "exception" def test_registration_injects_context_only_once() -> None: @@ -142,8 +158,8 @@ def test_registration_injects_context_only_once() -> None: def test_activate_and_deactivate_patch_all_restate_registration_paths( monkeypatch, ) -> None: - wrapped: list[tuple[str, str]] = [] - unwrapped: list[tuple[str, str]] = [] + installed = 0 + removed = 0 real_import = instrumentation.importlib.import_module def import_module(name: str): @@ -152,15 +168,38 @@ def import_module(name: str): return real_import(name) monkeypatch.setattr(instrumentation.importlib, "import_module", import_module) + + def install() -> None: + nonlocal installed + installed += 1 + + def remove() -> None: + nonlocal removed + removed += 1 + + monkeypatch.setattr(instrumentation, "_install_patches", install) + monkeypatch.setattr(instrumentation, "_remove_patches", remove) + monkeypatch.setattr(instrumentation, "_ACTIVATION_COUNT", 0) + monkeypatch.setattr(instrumentation, "_PATCHED_TARGETS", []) + monkeypatch.setattr(instrumentation, "_ENABLED", False) + + adapter = RestateInstrumentor() + adapter.activate() + assert installed == 1 + adapter.deactivate() + assert removed == 1 + + +def test_real_current_restate_registration_exports_connected_readable_span( + monkeypatch, +) -> None: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) monkeypatch.setattr( - instrumentation, - "wrap_function_wrapper", - lambda module, target, wrapper: wrapped.append((module, target)), - ) - monkeypatch.setattr( - instrumentation, - "unwrap", - lambda module, target: unwrapped.append((module, target)), + instrumentation.trace, + "get_tracer", + lambda *args, **kwargs: provider.get_tracer("restate", "0.1.0"), ) monkeypatch.setattr(instrumentation, "_ACTIVATION_COUNT", 0) monkeypatch.setattr(instrumentation, "_PATCHED_TARGETS", []) @@ -168,6 +207,97 @@ def import_module(name: str): adapter = RestateInstrumentor() adapter.activate() - assert wrapped == list(instrumentation.RESTATE_REGISTRATION_TARGETS) + workflow = restate.Workflow("CheckoutWorkflow", metadata={"team": "payments"}) + + @workflow.main(name="run") + async def run(_ctx, request: dict) -> dict: + return {"accepted": request["order_id"]} + + handler = workflow.handlers["run"] + assert instrumentation._invocation_context in handler.context_managers + + context = SimpleNamespace( + handler=handler, + invocation=SimpleNamespace( + invocation_id="inv-real-123", + input_buffer=b'{"order_id":"order-42"}', + key="order-42", + scope=None, + limit_key=None, + idempotency_key=None, + ), + ) + server_context = SimpleNamespace( + current_context=lambda: context, + restate_context_is_replaying=SimpleNamespace(get=lambda: False), + ) + real_import = instrumentation.importlib.import_module + monkeypatch.setattr( + instrumentation.importlib, + "import_module", + lambda name: ( + server_context if name == "restate.server_context" else real_import(name) + ), + ) + + async def invoke() -> None: + with provider.get_tracer("test").start_as_current_span("outer"): + async with instrumentation._invocation_context(): + pass + + asyncio.run(invoke()) adapter.deactivate() - assert unwrapped == list(reversed(instrumentation.RESTATE_REGISTRATION_TARGETS)) + + spans = exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "restate.workflow.CheckoutWorkflow.run", + "outer", + ] + child, root = spans + assert child.parent.span_id == root.context.span_id + assert child.status.status_code is StatusCode.OK + assert child.instrumentation_scope.name == "restate" + attrs = child.attributes + assert attrs[RESPAN_LOG_TYPE] == "workflow" + assert attrs[SpanAttributes.TRACELOOP_ENTITY_PATH] == "CheckoutWorkflow.run" + assert json.loads(attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT])["input"] == { + "order_id": "order-42" + } + assert "traceloop.span.kind" not in attrs + + +def test_serialization_is_bounded_redacted_and_hostile_safe() -> None: + class Hostile: + def __str__(self) -> str: + raise AssertionError("must not stringify") + + def __repr__(self) -> str: + raise AssertionError("must not repr") + + payload = { + "client_secret": "plain-secret", + "url": "https://user:password@example.com/restate?token=abc", + "emoji": "😀" * 5_000, + "hostile": Hostile(), + "nonfinite": float("nan"), + } + encoded = json_string(payload) + assert len(encoded.encode("utf-8")) <= 16_000 + assert "plain-secret" not in encoded + assert "password" not in encoded + assert "token=abc" not in encoded + assert json.loads(encoded)["nonfinite"] is None + + +def test_lifecycle_rejects_mismatched_capture_config(monkeypatch) -> None: + monkeypatch.setattr(instrumentation, "_ACTIVATION_COUNT", 0) + monkeypatch.setattr(instrumentation, "_PATCHED_TARGETS", []) + monkeypatch.setattr(instrumentation, "_install_patches", lambda: None) + monkeypatch.setattr(instrumentation, "_remove_patches", lambda: None) + + first = RestateInstrumentor(capture_content=True) + second = RestateInstrumentor(capture_content=False) + first.activate() + with pytest.raises(ValueError, match="capture_content"): + second.activate() + first.deactivate()