trodo-python 1.2.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

trodo/otel/auto_instrument.py

@@ -0,0 +1,281 @@
+"""Optional: register OpenTelemetry auto-instrumentation that feeds spans
+into our TrodoSpanProcessor. Call ``enable_auto_instrument()`` once at
+startup to pick up Anthropic / OpenAI / LangChain / LlamaIndex / Google
+Generative AI / Vertex AI and others — without any more code.
+
+Each underlying OTel package is an optional dependency — we import them
+lazily and skip anything that isn't installed.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, Iterable, List, Optional
+
+from .context import get_active_context
+from .processor import TrodoSpan, TrodoSpanProcessor
+
+
+def _hr_to_iso(nanos: Optional[int]) -> Optional[str]:
+    if nanos is None:
+        return None
+    return datetime.fromtimestamp(nanos / 1e9, tz=timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _infer_kind(attrs: Dict[str, Any]) -> str:
+    if not attrs:
+        return "generic"
+    if attrs.get("gen_ai.tool.name"):
+        return "tool"
+    if attrs.get("gen_ai.operation.name") or attrs.get("gen_ai.request.model"):
+        return "llm"
+    if attrs.get("db.system") or attrs.get("retrieval.query"):
+        return "retrieval"
+    return "generic"
+
+
+def otel_span_to_trodo_span(otel_span: Any) -> Optional[TrodoSpan]:
+    """Translate an OTel ReadableSpan to our TrodoSpan using GenAI semconv."""
+    ctx = get_active_context()
+    if ctx is None:
+        return None  # span emitted outside of wrap_agent — drop
+    try:
+        span_ctx = (
+            otel_span.get_span_context()
+            if callable(getattr(otel_span, "get_span_context", None))
+            else otel_span.context
+        )
+    except Exception:
+        return None
+
+    span_id_int = getattr(span_ctx, "span_id", None)
+    if not span_id_int:
+        return None
+    span_id = f"{span_id_int:016x}" if isinstance(span_id_int, int) else str(span_id_int)
+
+    parent = getattr(otel_span, "parent", None)
+    parent_span_id: Optional[str] = None
+    if parent is not None:
+        pid = getattr(parent, "span_id", None)
+        parent_span_id = f"{pid:016x}" if isinstance(pid, int) else str(pid) if pid else None
+    if parent_span_id is None:
+        parent_span_id = ctx.span_id
+
+    attrs = dict(getattr(otel_span, "attributes", {}) or {})
+    kind = _infer_kind(attrs)
+
+    start_time = getattr(otel_span, "start_time", None)
+    end_time = getattr(otel_span, "end_time", None)
+    started_at = _hr_to_iso(start_time) or datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+    ended_at = _hr_to_iso(end_time)
+    duration_ms = None
+    if start_time and end_time:
+        duration_ms = max(0, int((end_time - start_time) / 1e6))
+
+    status = getattr(otel_span, "status", None)
+    status_code = getattr(status, "status_code", None)
+    ok = "ok"
+    try:
+        from opentelemetry.trace import StatusCode
+
+        ok = "error" if status_code == StatusCode.ERROR else "ok"
+    except Exception:
+        ok = "error" if (status_code and str(status_code).endswith("ERROR")) else "ok"
+
+    # Accept both stable and experimental GenAI semconv keys.
+    in_toks = (
+        attrs.get("gen_ai.usage.input_tokens")
+        or attrs.get("gen_ai.usage.prompt_tokens")
+        or attrs.get("llm.usage.prompt_tokens")
+    )
+    out_toks = (
+        attrs.get("gen_ai.usage.output_tokens")
+        or attrs.get("gen_ai.usage.completion_tokens")
+        or attrs.get("llm.usage.completion_tokens")
+    )
+    model = (
+        attrs.get("gen_ai.request.model")
+        or attrs.get("gen_ai.response.model")
+        or attrs.get("llm.request.model")
+    )
+    provider = attrs.get("gen_ai.system") or attrs.get("llm.vendor")
+    prompt = attrs.get("gen_ai.prompt") or attrs.get("llm.prompts")
+    completion = attrs.get("gen_ai.completion") or attrs.get("llm.completion")
+
+    return TrodoSpan(
+        span_id=span_id,
+        run_id=ctx.run_id,
+        parent_span_id=parent_span_id,
+        kind=kind,
+        name=getattr(otel_span, "name", kind),
+        status=ok,
+        started_at=started_at,
+        ended_at=ended_at,
+        duration_ms=duration_ms,
+        model=model,
+        provider=provider,
+        input_tokens=int(in_toks) if in_toks is not None else None,
+        output_tokens=int(out_toks) if out_toks is not None else None,
+        temperature=attrs.get("gen_ai.request.temperature"),
+        tool_name=attrs.get("gen_ai.tool.name"),
+        input=prompt,
+        output=completion,
+        attributes=attrs or None,
+    )
+
+
+class _OtelAdapter:
+    """Adapter matching opentelemetry.sdk.trace.SpanProcessor interface."""
+
+    def __init__(self, processor: TrodoSpanProcessor) -> None:
+        self._processor = processor
+
+    def on_start(self, span: Any, parent_context: Any = None) -> None:
+        pass
+
+    def on_end(self, span: Any) -> None:
+        trodo_span = otel_span_to_trodo_span(span)
+        if trodo_span is not None:
+            self._processor.enqueue_span(trodo_span)
+
+    def shutdown(self) -> None:
+        self._processor.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30_000) -> bool:
+        self._processor.force_flush()
+        return True
+
+
+_INSTRUMENTORS: List[tuple[str, Callable[[], Any]]] = []
+
+
+def _register_instrumentors() -> None:
+    """Build the list of known instrumentor factories.
+
+    Each entry is ``(framework_id, factory)`` where ``factory`` performs a
+    lazy import and returns an ``instrument()`` call. Failures are swallowed
+    so missing optional deps never break user code.
+    """
+    global _INSTRUMENTORS
+    if _INSTRUMENTORS:
+        return
+
+    def _anthropic() -> Any:
+        from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor  # type: ignore
+
+        AnthropicInstrumentor().instrument()
+
+    def _openai() -> Any:
+        from opentelemetry.instrumentation.openai import OpenAIInstrumentor  # type: ignore
+
+        OpenAIInstrumentor().instrument()
+
+    def _openai_v2() -> Any:
+        from opentelemetry.instrumentation.openai_v2 import OpenAIInstrumentor  # type: ignore
+
+        OpenAIInstrumentor().instrument()
+
+    def _langchain() -> Any:
+        from opentelemetry.instrumentation.langchain import LangChainInstrumentor  # type: ignore
+
+        LangChainInstrumentor().instrument()
+
+    def _llama_index() -> Any:
+        from opentelemetry.instrumentation.llama_index import LlamaIndexInstrumentor  # type: ignore
+
+        LlamaIndexInstrumentor().instrument()
+
+    def _google_generativeai() -> Any:
+        from opentelemetry.instrumentation.google_generativeai import (  # type: ignore
+            GoogleGenerativeAIInstrumentor,
+        )
+
+        GoogleGenerativeAIInstrumentor().instrument()
+
+    def _vertexai() -> Any:
+        from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor  # type: ignore
+
+        VertexAIInstrumentor().instrument()
+
+    def _bedrock() -> Any:
+        from opentelemetry.instrumentation.bedrock import BedrockInstrumentor  # type: ignore
+
+        BedrockInstrumentor().instrument()
+
+    def _cohere() -> Any:
+        from opentelemetry.instrumentation.cohere import CohereInstrumentor  # type: ignore
+
+        CohereInstrumentor().instrument()
+
+    def _mistralai() -> Any:
+        from opentelemetry.instrumentation.mistralai import MistralAiInstrumentor  # type: ignore
+
+        MistralAiInstrumentor().instrument()
+
+    def _haystack() -> Any:
+        from opentelemetry.instrumentation.haystack import HaystackInstrumentor  # type: ignore
+
+        HaystackInstrumentor().instrument()
+
+    def _httpx() -> Any:
+        from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor  # type: ignore
+
+        HTTPXClientInstrumentor().instrument()
+
+    def _requests() -> Any:
+        from opentelemetry.instrumentation.requests import RequestsInstrumentor  # type: ignore
+
+        RequestsInstrumentor().instrument()
+
+    _INSTRUMENTORS = [
+        ("anthropic", _anthropic),
+        ("openai", _openai),
+        ("openai_v2", _openai_v2),
+        ("langchain", _langchain),
+        ("llama_index", _llama_index),
+        ("google_generativeai", _google_generativeai),
+        ("vertexai", _vertexai),
+        ("bedrock", _bedrock),
+        ("cohere", _cohere),
+        ("mistralai", _mistralai),
+        ("haystack", _haystack),
+        ("httpx", _httpx),
+        ("requests", _requests),
+    ]
+
+
+def enable_auto_instrument(
+    processor: TrodoSpanProcessor,
+    disable: Optional[Iterable[str]] = None,
+) -> List[str]:
+    """Register OTel auto-instrumentations for installed packages.
+
+    Returns the list of framework ids that were actually instrumented.
+    Skipped silently if opentelemetry-sdk or the per-framework
+    instrumentation package isn't installed.
+    """
+    disabled = set(disable or [])
+    try:
+        from opentelemetry import trace  # type: ignore
+        from opentelemetry.sdk.trace import TracerProvider  # type: ignore
+    except Exception:
+        return []
+
+    provider = trace.get_tracer_provider()
+    if not isinstance(provider, TracerProvider):
+        provider = TracerProvider()
+        trace.set_tracer_provider(provider)
+
+    provider.add_span_processor(_OtelAdapter(processor))
+
+    _register_instrumentors()
+    active: List[str] = []
+    for name, register in _INSTRUMENTORS:
+        if name in disabled:
+            continue
+        try:
+            register()
+            active.append(name)
+        except Exception:
+            continue
+    return active

trodo/otel/context.py

@@ -0,0 +1,44 @@
+"""Active run/span context using contextvars — auto-propagates across awaits
+and threading.local-style sync code.
+"""
+
+from __future__ import annotations
+
+import contextvars
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class ActiveSpanContext:
+    run_id: str
+    span_id: str
+    parent_span_id: Optional[str]
+    team_site_id: str
+    processor: object  # TrodoSpanProcessor — avoid circular import
+
+
+_active: contextvars.ContextVar[Optional[ActiveSpanContext]] = contextvars.ContextVar(
+    "trodo_active_span", default=None,
+)
+
+
+def get_active_context() -> Optional[ActiveSpanContext]:
+    return _active.get()
+
+
+class run_with_context:
+    """Context manager that sets the active span context for its duration."""
+
+    def __init__(self, ctx: ActiveSpanContext) -> None:
+        self._ctx = ctx
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> ActiveSpanContext:
+        self._token = _active.set(self._ctx)
+        return self._ctx
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        if self._token is not None:
+            _active.reset(self._token)
+            self._token = None