otelzero 0.2.0__py3-none-any.whl

otelzero/hooks/_langsmith.py

@@ -0,0 +1,504 @@
+"""LangChain and LangGraph through LangSmith's *own* OpenTelemetry export.
+
+LangChain and LangGraph trace themselves through LangSmith, and the
+``langsmith`` SDK can export those traces as OpenTelemetry spans
+(``langsmith[otel]``). Each run of a runnable, chain, graph node, model or tool
+becomes one span, named after the run (``LangGraph``, ``RunnableSequence``, an
+agent's name, ``model``, a tool's name). It carries ``langsmith.span.kind``,
+``gen_ai.operation.name`` (``chain`` / ``chat`` / ``execute_tool`` / ...),
+``gen_ai.request.model``, ``gen_ai.usage.*`` and ``langsmith.metadata.*``.
+On otelzero's provider they pass through :mod:`otelzero.translate` on export,
+which rewrites them into the current GenAI conventions: the root run becomes
+``invoke_agent {run}``, a model call ``chat {model}``, a tool call
+``execute_tool {tool}``, and run metadata (``thread_id`` / ``conversation_id``,
+``user_name`` / ``user_id``) becomes ``gen_ai.conversation.id`` / ``user.name``.
+If the application had already set its own global provider, they are exported
+through that one, untranslated.
+``install_langchain_monitoring`` and
+``install_langgraph_monitoring`` wrap nothing. Instead they:
+
+* configure otelzero's OTLP provider as the global one, which is where
+  LangSmith's exporter sends its spans;
+* turn the export on through LangSmith's own variables:
+  ``LANGSMITH_TRACING=true``, ``LANGSMITH_OTEL_ENABLED=true`` and, unless a
+  LangSmith API key is configured, ``LANGSMITH_OTEL_ONLY=true`` (so nothing is
+  sent to LangSmith's cloud). Unless content capture
+  (``OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT``) is opted in,
+  ``LANGSMITH_HIDE_INPUTS`` / ``LANGSMITH_HIDE_OUTPUTS`` are set too, because
+  LangSmith otherwise exports every prompt and completion as
+  ``gen_ai.prompt`` / ``gen_ai.completion``. Values the user already set are
+  left alone, and an uninstall removes only the ones this module set.
+* register a callback handler through LangChain's own
+  ``register_configure_hook``. It isn't a wrapper: LangChain adds it to every
+  run's callbacks the way it adds its own tracers. It sends the events and
+  records each run's identity (below).
+* attach the shared :class:`~otelzero.hooks._native.NativeRunProcessor`, with
+  events off, to stamp otelzero's identity on LangSmith's spans.
+
+**Why the events come from callbacks, not spans.** LangSmith builds its spans
+on a background thread, from the runs it has queued, once it flushes. A run
+that has finished by then is created and ended in one go, before the spans of
+its children exist. So the spans neither arrive live nor form a tree that is
+open while they are processed, and :func:`~otelzero.context.run_context` is
+not in scope on that thread. The callbacks fire in the caller's context, as
+the run happens. Each top-level run (the outermost runnable, chain, graph or
+agent a caller invoked) is one ``agent_start`` / ``agent_end`` /
+``agent_error``. The model, the tools it called and its token usage are
+collected from the model and tool runs under it. A graph's nodes and
+subgraphs are part of the run, not runs of their own. The session is titled
+from the run's input, as before.
+
+**Identity.** A run's events and its spans name the same thread and user.
+For the events, the identity is settled when the top-level run starts, in the
+caller's context: what the run names in its metadata (``conversation_id`` /
+``thread_id`` / ``session_id``, ``user_name`` / ``user_id`` / ``user``, the
+keys :mod:`otelzero.translate.langsmith` maps), else the active
+:func:`~otelzero.context.run_context`, else the client's defaults. A thread
+named in the metadata keeps the surrounding scope's user and title. A span
+can't be tied back to the run it records:
+LangSmith starts its root spans under a fresh random trace and span id, on its
+own thread, and sets no run id on them. So LangSmith's spans are stamped with
+the *client's* identity (``user=``, ``session=``, the session title) when they
+start, and the translator replaces it with the run's own wherever the run names
+one in its metadata: ``config={"metadata": {"thread_id": ..., "user_name":
+...}}``, or LangGraph's ``config={"configurable": {"thread_id": ...}}``, which
+LangGraph copies into the metadata.
+
+Both install functions share one handler and one processor, so installing
+both reports each run once, under the framework label of the latest call.
+"""
+
+import os
+import threading
+import traceback as _traceback
+from contextvars import ContextVar
+from typing import Any, Dict, Optional
+from uuid import UUID
+
+from otelzero import otel, semconv
+from otelzero.context import run_context
+from otelzero.hooks import _native
+from otelzero.hooks._usage import UsageAccumulator, extract_usage
+from otelzero.telemetry import TelemetryClient
+
+# One processor for every LangSmith span, whichever installer attached it.
+_PROCESSOR_KEY = "LangSmith"
+_SCOPES = ("langsmith",)
+
+# The env var LangChain checks to instantiate the callback handler.
+_HOOK_ENV = "OTELZERO_LANGCHAIN_CALLBACKS"
+
+_lock = threading.Lock()
+_telemetry: Optional[TelemetryClient] = None
+# top-level run id -> what its runs revealed; any run id -> its top-level run.
+_runs: Dict[UUID, dict] = {}
+_root_of: Dict[UUID, UUID] = {}
+_hook_registered = False
+# Environment values this module set, so uninstall can put them back.
+_set_env: dict = {}
+
+
+# --- events from LangChain's callbacks -----------------------------------------
+
+
+def _intent(value: Any) -> Optional[str]:
+    """Best-effort intent text from a top-level run's inputs, to title the session."""
+    if isinstance(value, str):
+        return value
+    if isinstance(value, dict):
+        for key in ("input", "question", "query", "prompt"):
+            found = value.get(key)
+            if isinstance(found, str) and found:
+                return found
+        value = value.get("messages")
+    if isinstance(value, (list, tuple)):
+        for message in value:
+            if isinstance(message, list):  # on_chat_model_start: [[messages]]
+                found = _intent(message)
+                if found:
+                    return found
+                continue
+            if isinstance(message, tuple) and len(message) == 2 and message[0] in ("user", "human"):
+                return message[1] if isinstance(message[1], str) else None
+            if isinstance(message, dict) and message.get("role") in ("user", "human"):
+                content = message.get("content")
+                return content if isinstance(content, str) else None
+            if getattr(message, "type", None) == "human":
+                content = getattr(message, "content", None)
+                return content if isinstance(content, str) else None
+    return None
+
+
+def _model_of(serialized: Any, metadata: Any, invocation_params: Any) -> Optional[str]:
+    for source, key in (
+        (metadata, "ls_model_name"),
+        (invocation_params, "model"),
+        (invocation_params, "model_name"),
+    ):
+        if isinstance(source, dict) and isinstance(source.get(key), str) and source[key]:
+            return source[key]
+    kwargs = serialized.get("kwargs") if isinstance(serialized, dict) else None
+    if isinstance(kwargs, dict):
+        for key in ("model", "model_name"):
+            if isinstance(kwargs.get(key), str) and kwargs[key]:
+                return kwargs[key]
+    return None
+
+
+def _run_name(serialized: Any, kwargs: dict, default: str) -> str:
+    name = kwargs.get("name")
+    if isinstance(name, str) and name:
+        return name
+    if isinstance(serialized, dict):
+        if isinstance(serialized.get("name"), str) and serialized["name"]:
+            return serialized["name"]
+        ids = serialized.get("id")
+        if isinstance(ids, list) and ids and isinstance(ids[-1], str):
+            return ids[-1]
+    return default
+
+
+def _usage_of(response: Any) -> Optional[dict]:
+    """Token usage of one model call, from its ``LLMResult``."""
+    for generations in getattr(response, "generations", None) or []:
+        for generation in generations or []:
+            message = getattr(generation, "message", None)
+            usage = extract_usage(getattr(message, "usage_metadata", None))
+            if usage:
+                return usage
+    llm_output = getattr(response, "llm_output", None)
+    if isinstance(llm_output, dict):
+        return extract_usage(llm_output.get("token_usage") or llm_output.get("usage"))
+    return None
+
+
+def _metadata_identity(metadata: Any):
+    """``(user, conversation_id)`` a run names in its metadata, else ``None``s.
+
+    The same keys, in the same order, that :mod:`otelzero.translate.langsmith`
+    maps onto the run's spans, so a run's events and spans agree. LangGraph
+    copies ``config["configurable"]["thread_id"]`` into the metadata, so a
+    LangGraph thread counts too. Imported lazily: ``otelzero.translate`` needs
+    the OpenTelemetry SDK, and this module is imported with ``otelzero``.
+    """
+    if not isinstance(metadata, dict) or not metadata:
+        return None, None
+    from otelzero.translate.langsmith import CONVERSATION_KEYS, USER_KEYS
+
+    def first(keys):
+        for key in keys:
+            value = metadata.get(key)
+            if value is not None and value != "":
+                return str(value)
+        return None
+
+    return first(USER_KEYS), first(CONVERSATION_KEYS)
+
+
+def _send(telemetry: TelemetryClient, run: dict, **fields: Any) -> None:
+    """Send one of a run's events under the identity it started with."""
+    user, conversation_id, session_name = run["identity"]
+    with run_context(user=user, conversation_id=conversation_id, name=session_name):
+        telemetry.send_log(**fields)
+
+
+def _start(
+    run_id: Any,
+    parent_run_id: Any,
+    name: str,
+    inputs: Any,
+    model: Optional[str],
+    metadata: Any = None,
+) -> None:
+    """Register a run; a top-level one also sends ``agent_start``.
+
+    The run's identity is settled here, in the caller's context: what the run
+    names in its metadata, else the active ``run_context``, else the client's
+    defaults. All of the run's events carry it.
+    """
+    telemetry = _telemetry
+    if telemetry is None or not isinstance(run_id, UUID):
+        return
+    with _lock:
+        if parent_run_id is not None:
+            root = _root_of.get(parent_run_id)
+            if root is not None:
+                _root_of[run_id] = root
+                run = _runs.get(root)
+                if run is not None and model and not run["model"]:
+                    run["model"] = model
+            return
+        _root_of[run_id] = run_id
+        run = _runs[run_id] = {
+            "name": name,
+            "model": model,
+            "tools": set(),
+            "usage": UsageAccumulator(),
+        }
+    meta_user, meta_conversation = _metadata_identity(metadata)
+    user, conversation_id, session_name = telemetry.run_identity()
+    if meta_conversation is not None:
+        # The run's own thread, inside whatever scope it was started in: keep
+        # that scope's user and title (an explicit ``name=`` included), and
+        # title the thread from this run's input only if there is none yet.
+        with run_context(user=user, conversation_id=meta_conversation, name=session_name):
+            telemetry.suggest_session_name(_intent(inputs))
+            _, _, session_name = telemetry.run_identity()
+        conversation_id = meta_conversation
+    else:
+        telemetry.suggest_session_name(_intent(inputs))
+        user, conversation_id, session_name = telemetry.run_identity()
+    run["identity"] = (meta_user or user, conversation_id, session_name)
+    _send(
+        telemetry,
+        run,
+        event_type="agent_start",
+        agent_name=name,
+        agent_role=name,
+        model=model,
+        details={"model": model},
+    )
+
+
+def _end(run_id: Any, error: Optional[BaseException] = None) -> None:
+    """Forget a run; a top-level one also sends ``agent_end`` / ``agent_error``."""
+    with _lock:
+        root = _root_of.pop(run_id, None)
+        run = _runs.pop(run_id, None) if root == run_id else None
+        if run is not None:
+            for child in [k for k, v in _root_of.items() if v == run_id]:
+                del _root_of[child]
+    telemetry = _telemetry
+    if run is None or telemetry is None:
+        return
+    fields = {"agent_name": run["name"], "agent_role": run["name"], "model": run["model"]}
+    usage = run["usage"].result()
+    if error is not None:
+        _send(
+            telemetry,
+            run,
+            event_type="agent_error",
+            error=str(error) or type(error).__name__,
+            traceback="".join(
+                _traceback.format_exception(type(error), error, error.__traceback__)
+            ),
+            details={"usage": usage} if usage else None,
+            **fields,
+        )
+        return
+    details: dict = {"status": "success"}
+    if run["tools"]:
+        details["tools"] = sorted(run["tools"])
+    if usage:
+        details["usage"] = usage
+    _send(telemetry, run, event_type="agent_end", details=details, **fields)
+
+
+def _root_run(run_id: Any) -> Optional[dict]:
+    with _lock:
+        root = _root_of.get(run_id)
+        return _runs.get(root) if root is not None else None
+
+
+def _handler_class():
+    from langchain_core.callbacks import BaseCallbackHandler
+
+    class OtelzeroRunEvents(BaseCallbackHandler):
+        """Sends otelzero's ``agent_*`` events for top-level runs; changes nothing."""
+
+        # Inline, so the sync callbacks run in the caller's context (where
+        # run_context is) even under an async run.
+        run_inline = True
+        raise_error = False
+
+        # -- starts ---------------------------------------------------------
+
+        def on_chain_start(
+            self, serialized, inputs, *, run_id, parent_run_id=None, metadata=None, **kwargs
+        ):
+            name = _run_name(serialized, kwargs, "chain")
+            _start(run_id, parent_run_id, name, inputs, None, metadata)
+
+        def on_chat_model_start(
+            self, serialized, messages, *, run_id, parent_run_id=None, metadata=None, **kwargs
+        ):
+            model = _model_of(serialized, metadata, kwargs.get("invocation_params"))
+            name = _run_name(serialized, kwargs, "chat_model")
+            _start(run_id, parent_run_id, name, messages, model, metadata)
+
+        def on_llm_start(
+            self, serialized, prompts, *, run_id, parent_run_id=None, metadata=None, **kwargs
+        ):
+            model = _model_of(serialized, metadata, kwargs.get("invocation_params"))
+            name = _run_name(serialized, kwargs, "llm")
+            _start(run_id, parent_run_id, name, prompts, model, metadata)
+
+        def on_tool_start(
+            self, serialized, input_str, *, run_id, parent_run_id=None, metadata=None, **kwargs
+        ):
+            name = _run_name(serialized, kwargs, "tool")
+            _start(run_id, parent_run_id, name, input_str, None, metadata)
+            run = _root_run(run_id)
+            if run is not None:
+                with _lock:
+                    run["tools"].add(name)
+
+        def on_retriever_start(
+            self, serialized, query, *, run_id, parent_run_id=None, metadata=None, **kwargs
+        ):
+            name = _run_name(serialized, kwargs, "retriever")
+            _start(run_id, parent_run_id, name, query, None, metadata)
+
+        # -- ends -----------------------------------------------------------
+
+        def on_llm_end(self, response, *, run_id, **kwargs):
+            run = _root_run(run_id)
+            if run is not None:
+                usage = _usage_of(response)
+                with _lock:
+                    run["usage"].add(usage)
+            _end(run_id)
+
+        def on_chain_end(self, outputs, *, run_id, **kwargs):
+            _end(run_id)
+
+        def on_tool_end(self, output, *, run_id, **kwargs):
+            _end(run_id)
+
+        def on_retriever_end(self, documents, *, run_id, **kwargs):
+            _end(run_id)
+
+        def on_chain_error(self, error, *, run_id, **kwargs):
+            _end(run_id, error)
+
+        def on_llm_error(self, error, *, run_id, **kwargs):
+            _end(run_id, error)
+
+        def on_tool_error(self, error, *, run_id, **kwargs):
+            _end(run_id, error)
+
+        def on_retriever_error(self, error, *, run_id, **kwargs):
+            _end(run_id, error)
+
+    return OtelzeroRunEvents
+
+
+def _register_handler() -> None:
+    global _hook_registered
+    if _hook_registered:
+        return
+    from langchain_core.tracers.context import register_configure_hook
+
+    register_configure_hook(
+        ContextVar("otelzero_run_events", default=None),
+        True,
+        _handler_class(),
+        _HOOK_ENV,
+    )
+    _hook_registered = True
+
+
+# --- LangSmith's export ------------------------------------------------------
+
+
+def _setdefault_env(name: str, value: str) -> None:
+    if name in os.environ:
+        return
+    os.environ[name] = value
+    _set_env[name] = value
+
+
+def _enable_export() -> None:
+    _setdefault_env("LANGSMITH_TRACING", "true")
+    _setdefault_env("LANGSMITH_OTEL_ENABLED", "true")
+    if not (os.environ.get("LANGSMITH_API_KEY") or os.environ.get("LANGCHAIN_API_KEY")):
+        # Nothing to send to LangSmith's cloud: export over OTel only.
+        _setdefault_env("LANGSMITH_OTEL_ONLY", "true")
+    if not semconv.capture_content_enabled():
+        _setdefault_env("LANGSMITH_HIDE_INPUTS", "true")
+        _setdefault_env("LANGSMITH_HIDE_OUTPUTS", "true")
+    os.environ[_HOOK_ENV] = "true"
+    _clear_langsmith_env_cache()
+
+
+def _clear_langsmith_env_cache() -> None:
+    """LangSmith caches env lookups; drop them so the new values are read."""
+    try:
+        from langsmith import utils as ls_utils
+    except ImportError:  # pragma: no cover - checked by install
+        return
+    for name in ("get_env_var", "tracing_is_enabled", "get_tracer_project"):
+        cache_clear = getattr(getattr(ls_utils, name, None), "cache_clear", None)
+        if callable(cache_clear):
+            cache_clear()
+
+
+def install(
+    framework: str,
+    server_url: str,
+    user: Optional[str],
+    api_key: Optional[str],
+    session: Optional[str],
+    name: Optional[str],
+) -> TelemetryClient:
+    """Shared body of ``install_langchain_monitoring`` / ``install_langgraph_monitoring``."""
+    global _telemetry
+    try:
+        import langchain_core  # noqa: F401
+        import langsmith  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - exercised by hand
+        raise ImportError(
+            f"{framework} monitoring needs langchain-core and langsmith importable: "
+            "uv pip install langchain-core 'langsmith[otel]'"
+        ) from exc
+
+    telemetry = TelemetryClient(
+        server_url,
+        framework=framework,
+        user=user,
+        api_key=api_key,
+        session=session,
+        name=name,
+    )
+    provider = _native.start_tracing(framework, telemetry, server_url)
+    _native.attach_processor(
+        _native.global_provider(provider, framework),
+        telemetry,
+        framework=framework,
+        scopes=_SCOPES,
+        default_agent_name=f"{framework} run",
+        send_events=False,
+        key=_PROCESSOR_KEY,
+    )
+    _telemetry = telemetry
+    _enable_export()
+    _register_handler()
+
+    print(
+        f"[MONITOR] {framework} monitoring installed via LangSmith's native "
+        f"OpenTelemetry export (server: {server_url})"
+    )
+    return telemetry
+
+
+def uninstall() -> None:
+    """Stop the events and span stamping; undo the env vars this module set.
+
+    The callback handler stays registered (LangChain offers no way to remove a
+    configure hook) but is no longer instantiated, and is dormant if it is.
+    """
+    global _telemetry
+    from opentelemetry import trace
+
+    _native.detach(_PROCESSOR_KEY, otel._provider)
+    _native.detach(_PROCESSOR_KEY, trace.get_tracer_provider())
+    _telemetry = None
+    os.environ.pop(_HOOK_ENV, None)
+    for env_name, value in list(_set_env.items()):
+        if os.environ.get(env_name) == value:
+            del os.environ[env_name]
+        del _set_env[env_name]
+    with _lock:
+        _runs.clear()
+        _root_of.clear()
+    _clear_langsmith_env_cache()