otelzero 0.2.0__py3-none-any.whl

otelzero/hooks/databricks.py

@@ -0,0 +1,260 @@
+"""Databricks Agent Bricks / Mosaic AI instrumentation for MLflow agents."""
+
+from typing import Optional
+
+import wrapt
+
+from otelzero import semconv
+from otelzero.hooks import _content
+from otelzero.hooks._usage import _field, extract_usage
+from otelzero.hooks.base import _ImportInstrumenter, install_hook, trace_agent_run
+
+# Cap free-text task detail (tool arguments / output) so a large payload can't
+# bloat every task event; mirrors the [:100] truncation the other hooks use.
+_MAX_DETAIL = 200
+
+
+def _truncate(value):
+    """Coerce to a short string for a telemetry ``details`` field (or ``None``)."""
+    if value is None:
+        return None
+    text = value if isinstance(value, str) else str(value)
+    return text if len(text) <= _MAX_DETAIL else text[:_MAX_DETAIL] + "…"
+
+
+def _first_present(obj, *keys):
+    """First of ``keys`` that reads as a non-``None`` field off ``obj``."""
+    for key in keys:
+        value = _field(obj, key)
+        if value is not None:
+            return value
+    return None
+
+
+def _emit_tool_task_events(telemetry, agent_name, response):
+    """Surface each tool a ResponsesAgent called as a otelzero task on its timeline.
+
+    A ``ResponsesAgent`` records its tool calls in ``response.output`` as
+    ``function_call`` items, with results in matching ``function_call_output``
+    items (paired by ``call_id``). Each call is a unit of work the run did, so we
+    emit a task_start/task_end pair for it -- retrospectively, since the tool has
+    already executed inside ``predict`` (the ``with`` body is empty; we're
+    recording the completed call, not wrapping a live one). ``agent_role`` is set
+    to the agent name so the task attaches to that agent's card rather than
+    spawning its own (see ``kiji_agent_registry.db``'s task-event guard).
+
+    Responses without an ``output`` list -- a plain-text answer, or a ``ChatAgent``
+    (whose response carries ``messages``, not ``output``) -- yield no tasks.
+    """
+    output = _field(response, "output")
+    if not isinstance(output, (list, tuple)) or not output:
+        return
+
+    # Index tool results by call_id so each task_end can report its outcome.
+    results = {
+        _field(it, "call_id"): it
+        for it in output
+        if _field(it, "type") == "function_call_output"
+    }
+
+    for item in output:
+        if _field(item, "type") != "function_call":
+            continue
+        tool = _field(item, "name") or "tool"
+        call_id = _field(item, "call_id")
+        start_details = {
+            "tool": tool,
+            "call_id": call_id,
+            "arguments": _truncate(_field(item, "arguments")),
+        }
+        end_details = {"status": "success"}
+        result = results.get(call_id)
+        if result is not None:
+            end_details["output"] = _truncate(_field(result, "output"))
+
+        with trace_agent_run(
+            telemetry,
+            semconv.span_name(semconv.OP_EXECUTE_TOOL, tool),
+            {
+                semconv.GEN_AI_OPERATION_NAME: semconv.OP_EXECUTE_TOOL,
+                semconv.GEN_AI_AGENT_NAME: agent_name,
+                semconv.GEN_AI_TOOL_NAME: tool,
+                semconv.GEN_AI_TOOL_CALL_ID: call_id,
+            },
+            start_event="task_start",
+            end_event="task_end",
+            error_event="task_error",
+            log_fields={"agent_name": agent_name, "agent_role": agent_name},
+            start_details=start_details,
+            end_details=end_details,
+        ):
+            pass
+
+
+class DatabricksImportHook(_ImportInstrumenter):
+    """Instrument Databricks Agent Bricks / Mosaic AI agents.
+
+    Databricks agents (Agent Bricks, the Mosaic AI Agent Framework) are authored
+    as MLflow agent subclasses -- ``mlflow.pyfunc.ResponsesAgent`` (the current
+    recommended interface) or ``mlflow.pyfunc.ChatAgent`` (its predecessor).
+    Both live in ``mlflow.pyfunc.model``, both declare ``predict`` as an
+    *abstract* method each agent implements, and both drive subclass validation
+    through ``__init_subclass__`` -- so, unlike a concrete framework class, there
+    is no single method on the base to wrap. Instead this hook intercepts
+    subclass creation (``__init_subclass__``) for each base and wraps each
+    concrete agent's ``predict`` as it is defined (and any subclasses already
+    defined when the module loads), composing on top of MLflow's own validation
+    wrapper. Calling ``agent.predict(...)`` -- directly or via MLflow serving --
+    then streams agent telemetry with zero changes to the agent code.
+
+    Each ``predict`` call emits agent_start/agent_end/agent_error. On top of that,
+    the tools a ``ResponsesAgent`` invokes -- recorded in its response as
+    ``function_call`` items (with results in matching ``function_call_output``
+    items) -- are surfaced as nested task_start/task_end events, one per tool, so
+    a run's tool steps show up on the agent's timeline (see
+    :func:`_emit_tool_task_events`).
+
+    Note: only ``predict`` is wrapped (the required entry point). ``predict_stream``
+    is left alone to avoid double-counting agents whose ``predict`` aggregates
+    ``predict_stream`` internally.
+    """
+
+    WATCHED = {"mlflow.pyfunc.model": "_wrap_agent_bases"}
+
+    # otelzero's own marker so we never double-wrap (MLflow already sets __wrapped__).
+    _MARKER = "_otelzero_wrapped"
+
+    # The MLflow agent-authoring bases to instrument, both in mlflow.pyfunc.model.
+    _AGENT_BASES = ("ResponsesAgent", "ChatAgent")
+
+    def _wrap_agent_bases(self, module):
+        for base_name in self._AGENT_BASES:
+            agent_base = getattr(module, base_name, None)
+            if agent_base is not None:
+                self._instrument_base(agent_base)
+
+    def _instrument_base(self, agent_base):
+        telemetry = self.telemetry
+
+        @wrapt.decorator
+        def predict_wrapper(wrapped, instance, args, kwargs):
+            agent_name = type(instance).__name__
+            # ChatAgent.predict(messages, ...) receives the message list
+            # directly; ResponsesAgent.predict(request) wraps it in
+            # ``request.input`` -- unwrap that so both count correctly.
+            arg = args[0] if args else kwargs.get("messages", kwargs.get("request"))
+            payload = getattr(arg, "input", arg)
+            try:
+                num_messages = len(payload)
+            except TypeError:
+                num_messages = 0
+            agent_info = {"num_messages": num_messages}
+
+            with trace_agent_run(
+                telemetry,
+                semconv.span_name(semconv.OP_INVOKE_AGENT, agent_name),
+                {
+                    semconv.GEN_AI_OPERATION_NAME: semconv.OP_INVOKE_AGENT,
+                    semconv.GEN_AI_AGENT_NAME: agent_name,
+                    semconv.OTELZERO_NUM_MESSAGES: num_messages,
+                },
+                start_event="agent_start",
+                end_event="agent_end",
+                error_event="agent_error",
+                log_fields={"agent_name": agent_name, "agent_role": agent_name},
+                start_details=agent_info,
+            ) as run:
+                # ChatAgent messages ({role, content}) or ResponsesAgent input
+                # items (message / function_call / function_call_output).
+                run.set_input(_content.messages_from(payload))
+                response = wrapped(*args, **kwargs)
+                # Surface the tools this run called as nested task events. Emitted
+                # inside the agent span (so they attach to its trace/timeline)
+                # and after predict returns (the calls have already run).
+                _emit_tool_task_events(telemetry, agent_name, response)
+                # Responses carry OpenAI-style usage (input/output tokens plus
+                # ``*_tokens_details.reasoning_tokens`` when reported).
+                run.set_usage(extract_usage(_field(response, "usage")))
+                # ChatAgentResponse.messages, or ResponsesAgentResponse.output
+                # items (the same message / tool item shapes as the input).
+                run.set_output(
+                    _content.messages_from(
+                        _first_present(response, "messages", "output"),
+                        default_role="assistant",
+                    )
+                    or _content.text_of(response)
+                )
+                return response
+
+        marker = self._MARKER
+
+        def _wrap_predict(cls):
+            predict = cls.__dict__.get("predict")
+            if (
+                predict is None
+                or getattr(predict, "__isabstractmethod__", False)
+                or getattr(predict, marker, False)
+            ):
+                return
+            wrapped = predict_wrapper(predict)
+            setattr(wrapped, marker, True)
+            cls.predict = wrapped
+
+        # Wrap concrete subclasses that already exist when the module loads ...
+        def _walk(cls):
+            for sub in cls.__subclasses__():
+                _wrap_predict(sub)
+                _walk(sub)
+
+        _walk(agent_base)
+
+        # ... and intercept every future subclass definition. ``__init_subclass__``
+        # is implicitly a classmethod; chain to the previous one (MLflow installs
+        # its data-validation wrapper there) before applying ours on top.
+        prev_init_subclass = agent_base.__dict__.get("__init_subclass__")
+
+        def _patched_init_subclass(cls, **kwargs):
+            if prev_init_subclass is not None:
+                prev_init_subclass.__func__(cls, **kwargs)
+            else:
+                super(agent_base, cls).__init_subclass__(**kwargs)
+            _wrap_predict(cls)
+
+        agent_base.__init_subclass__ = classmethod(_patched_init_subclass)
+
+
+def install_databricks_monitoring(
+    server_url: str = "http://localhost:8000",
+    user: Optional[str] = None,
+    api_key: Optional[str] = None,
+    session: Optional[str] = None,
+    name: Optional[str] = None,
+):
+    """
+    Install Databricks Agent Bricks / Mosaic AI monitoring instrumentation.
+
+    Instruments MLflow ``ResponsesAgent`` and ``ChatAgent`` subclasses -- the
+    authoring interfaces for Databricks Agent Bricks and the Mosaic AI Agent
+    Framework -- by wrapping each concrete agent's ``predict`` as it is defined.
+    Every ``agent.predict(...)`` call (direct or via MLflow serving) streams
+    agent_start/agent_end/agent_error telemetry with zero changes to the agent
+    code. Call this *before* ``mlflow.pyfunc`` is imported (i.e. before defining
+    your agent subclass).
+
+    Args:
+        server_url: URL of the monitoring server
+        user: Name of the executing user
+        api_key: Ingest key tying events to an owner (for the iOS app)
+        session: Per-run id used to group events in the app
+        name: Explicit session title (else falls back to a codename)
+    """
+    return install_hook(
+        DatabricksImportHook,
+        framework="Databricks Agents",
+        label="Databricks Agent Bricks",
+        server_url=server_url,
+        user=user,
+        api_key=api_key,
+        session=session,
+        name=name,
+    )

otelzero/hooks/langchain.py

@@ -0,0 +1,56 @@
+"""LangChain instrumentation: switch on LangSmith's *own* OpenTelemetry export.
+
+LangChain traces itself through LangSmith, and ``langsmith[otel]`` exports
+those traces as OpenTelemetry spans. Every chain, runnable, agent, model and
+tool call becomes a span named after the run, carrying ``langsmith.span.kind``,
+``gen_ai.operation.name``, ``gen_ai.request.model`` and ``gen_ai.usage.*``.
+That is more than wrapping ``invoke`` from the outside could see, so this hook
+wraps nothing: it points the export at the registry and derives otelzero's
+events from it. See :mod:`otelzero.hooks._langsmith` for the details, shared
+with :mod:`otelzero.hooks.langgraph`.
+"""
+
+from typing import Optional
+
+from otelzero.hooks import _langsmith
+from otelzero.telemetry import TelemetryClient
+
+FRAMEWORK = "LangChain"
+
+
+def install_langchain_monitoring(
+    server_url: str = "http://localhost:8000",
+    user: Optional[str] = None,
+    api_key: Optional[str] = None,
+    session: Optional[str] = None,
+    name: Optional[str] = None,
+) -> TelemetryClient:
+    """
+    Install LangChain monitoring via LangSmith's native OpenTelemetry export.
+
+    Exports LangChain's own LangSmith spans to ``{server_url}/v1/traces`` (or
+    to ``OTEL_EXPORTER_OTLP_ENDPOINT`` when set), translated into the current
+    GenAI conventions on export (:mod:`otelzero.translate`), and
+    sends ``agent_start`` and ``agent_end`` (or ``agent_error``) for each
+    top-level run: a chain, a runnable, or a LangChain 1.0 ``create_agent``
+    agent (a LangGraph graph). May be called before or after ``langchain`` is
+    imported; requires ``langsmith[otel]`` and the ``otel`` extra.
+
+    Args:
+        server_url: URL of the monitoring server
+        user: Name of the executing user
+        api_key: Ingest key tying events to an owner (for the iOS app)
+        session: Per-run id used to group events in the app; exported on every
+            span as ``gen_ai.conversation.id``
+        name: Explicit session title (else derived from the run's input)
+
+    Raises:
+        ImportError: ``langchain-core`` or ``langsmith`` is not installed.
+        RuntimeError: the OpenTelemetry SDK (the ``otel`` extra) is missing.
+    """
+    return _langsmith.install(FRAMEWORK, server_url, user, api_key, session, name)
+
+
+def uninstall_langchain_monitoring() -> None:
+    """Stop the events; undo the LangSmith env vars the install set."""
+    _langsmith.uninstall()

otelzero/hooks/langgraph.py

@@ -0,0 +1,57 @@
+"""LangGraph instrumentation: switch on LangSmith's *own* OpenTelemetry export.
+
+LangGraph traces itself through LangSmith like the rest of LangChain: a graph
+run is a LangSmith run, each node and subgraph a child run, and every model and
+tool call a run below those. With ``langsmith[otel]`` each becomes an
+OpenTelemetry span. LangChain 1.0's ``create_agent`` and ``create_react_agent``
+return such graphs. This hook wraps nothing: it points the export at the
+registry and turns each graph run (the trace's root span) into
+``agent_start`` / ``agent_end`` / ``agent_error``. Its nodes and subgraphs are
+part of the run, not runs of their own. See :mod:`otelzero.hooks._langsmith`
+for the details, shared with :mod:`otelzero.hooks.langchain`.
+"""
+
+from typing import Optional
+
+from otelzero.hooks import _langsmith
+from otelzero.telemetry import TelemetryClient
+
+FRAMEWORK = "LangGraph"
+
+
+def install_langgraph_monitoring(
+    server_url: str = "http://localhost:8000",
+    user: Optional[str] = None,
+    api_key: Optional[str] = None,
+    session: Optional[str] = None,
+    name: Optional[str] = None,
+) -> TelemetryClient:
+    """
+    Install LangGraph monitoring via LangSmith's native OpenTelemetry export.
+
+    Exports LangGraph's own LangSmith spans to ``{server_url}/v1/traces`` (or
+    to ``OTEL_EXPORTER_OTLP_ENDPOINT`` when set), translated into the current
+    GenAI conventions on export (:mod:`otelzero.translate`), and
+    sends ``agent_start`` and ``agent_end`` (or ``agent_error``) for each graph
+    run. May be called before or after ``langgraph`` is imported; requires
+    ``langsmith[otel]`` and the ``otel`` extra. Installing this and
+    ``install_langchain_monitoring`` together still reports each run once.
+
+    Args:
+        server_url: URL of the monitoring server
+        user: Name of the executing user
+        api_key: Ingest key tying events to an owner (for the iOS app)
+        session: Per-run id used to group events in the app; exported on every
+            span as ``gen_ai.conversation.id``
+        name: Explicit session title (else derived from the graph's input)
+
+    Raises:
+        ImportError: ``langchain-core`` or ``langsmith`` is not installed.
+        RuntimeError: the OpenTelemetry SDK (the ``otel`` extra) is missing.
+    """
+    return _langsmith.install(FRAMEWORK, server_url, user, api_key, session, name)
+
+
+def uninstall_langgraph_monitoring() -> None:
+    """Stop the events; undo the LangSmith env vars the install set."""
+    _langsmith.uninstall()