otelzero 0.2.0__py3-none-any.whl

otelzero/hooks/claude_sdk.py

@@ -0,0 +1,267 @@
+"""Claude Agent SDK instrumentation: wrap the async-generator run chokepoints."""
+
+from typing import Any, Optional
+
+import wrapt
+
+from otelzero import semconv
+from otelzero.hooks import _content
+from otelzero.hooks._usage import UsageAccumulator, extract_usage
+from otelzero.hooks.base import _ImportInstrumenter, install_hook, trace_agent_run
+
+
+def _claude_model_id(options: Any) -> Optional[str]:
+    """Best-effort model identifier for a ``ClaudeAgentOptions`` object.
+
+    ``ClaudeAgentOptions.model`` is an optional model string (``None`` selects
+    the SDK/CLI default). Returns the string when set, else ``None``.
+    """
+    if options is None:
+        return None
+    model = getattr(options, "model", None)
+    if isinstance(model, str) and model:
+        return model
+    return None
+
+
+def _options_details(options: Any) -> dict:
+    """Summarise a ``ClaudeAgentOptions`` for a telemetry ``details`` payload.
+
+    Every field read degrades to a benign default so a run with no options (or
+    an SDK whose options shape has drifted) never breaks instrumentation.
+    """
+    allowed = getattr(options, "allowed_tools", None)
+    mcp = getattr(options, "mcp_servers", None)
+    return {
+        "model": _claude_model_id(options),
+        "allowed_tools": list(allowed) if allowed else [],
+        # ClaudeAgentOptions.mcp_servers is a name -> config dict; surface just
+        # the server names (configs can hold secrets and are large).
+        "mcp_servers": sorted(mcp.keys()) if isinstance(mcp, dict) else None,
+    }
+
+
+class ClaudeAgentSDKImportHook(_ImportInstrumenter):
+    """Intercept Claude Agent SDK imports and instrument both run chokepoints.
+
+    The SDK exposes two entry points, both async generators, and this hook wraps
+    each so a run is captured exactly once whichever style the caller uses:
+
+    * ``claude_agent_sdk.query(prompt, options)`` -- the one-shot path, a
+      module-level async-generator *function* (so the wrapper's ``instance`` is
+      ``None`` and options come from ``kwargs``).
+    * ``ClaudeSDKClient.receive_response()`` -- the interactive path, an
+      async-generator *method* that yields one response turn's messages (options
+      live on the client instance).
+
+    ``ClaudeSDKClient.query()`` is wrapped as well, though it is not a run: it
+    is where the interactive path's prompt goes out, so the prompt is recorded
+    there and attached to the ``receive_response()`` run that follows.
+
+    All are re-exported from the ``claude_agent_sdk`` package ``__init__``, so a
+    single watched module covers them. The SDK's own ``hooks=`` option
+    (PreToolUse/PostToolUse) is unrelated to otelzero's import hooks.
+    """
+
+    WATCHED = {"claude_agent_sdk": "_wrap_module"}
+
+    # otelzero's own marker so we never double-wrap a chokepoint.
+    _MARKER = "_otelzero_wrapped"
+
+    def _wrap_module(self, module):
+        self._wrap_query(module)
+        self._wrap_client(module)
+
+    def _wrap_query(self, module):
+        """Wrap the module-level ``query`` async-generator function."""
+        query = getattr(module, "query", None)
+        if query is None or getattr(query, self._MARKER, False):
+            return
+
+        telemetry = self.telemetry
+
+        @wrapt.decorator
+        async def query_wrapper(wrapped, instance, args, kwargs):
+            # ``query`` takes options by keyword (``query(prompt, options=...)``).
+            options = kwargs.get("options")
+            details = _options_details(options)
+            # The prompt is the run's intent. Only a string names the session
+            # (the streaming-message form has no single title). As input, a
+            # string or a materialised list is recorded up front; a streamed
+            # prompt is the SDK's to read, so the SDK gets a tee of the
+            # caller's iterable and the messages are recorded as they flow
+            # past (see ``_content.PromptTee``).
+            prompt = kwargs.get("prompt", args[0] if args else None)
+            prompt_text = prompt if isinstance(prompt, str) else None
+            tee = _content.PromptTee.wrap(prompt)
+            if tee is not None:
+                args, kwargs = _content.replace_prompt(args, kwargs, tee)
+
+            with trace_agent_run(
+                telemetry,
+                semconv.span_name(semconv.OP_INVOKE_AGENT, "Claude Agent"),
+                {
+                    semconv.GEN_AI_OPERATION_NAME: semconv.OP_INVOKE_AGENT,
+                    semconv.GEN_AI_AGENT_NAME: "Claude Agent",
+                    semconv.GEN_AI_REQUEST_MODEL: details["model"],
+                    semconv.GEN_AI_TOOL_DEFINITIONS: semconv.tool_definitions(
+                        details["allowed_tools"]
+                    ),
+                },
+                start_event="agent_start",
+                end_event="agent_end",
+                error_event="agent_error",
+                log_fields={
+                    "agent_name": "Claude Agent",
+                    "agent_role": "Claude Agent",
+                    "model": details["model"],
+                },
+                start_details=details,
+                session_name_hint=prompt_text,
+            ) as run:
+                run.set_input(_content.prompt_messages(prompt))
+                # Only the final ResultMessage carries ``.usage`` (input/output/
+                # cache-read token counts); other message types yield None.
+                # AssistantMessage/UserMessage content blocks (text, tool use,
+                # tool results) accumulate into the run's output messages.
+                acc = UsageAccumulator()
+                out = _content.OutputAccumulator()
+                try:
+                    async for message in wrapped(*args, **kwargs):
+                        acc.add(extract_usage(getattr(message, "usage", None)))
+                        out.add(message)
+                        yield message
+                finally:
+                    if tee is not None:
+                        run.set_input(tee.result())
+                    run.set_usage(acc.result())
+                    run.set_output(out.result(), finish_reasons=out.finish_reasons)
+
+        wrapped = query_wrapper(query)
+        setattr(wrapped, self._MARKER, True)
+        module.query = wrapped
+
+    def _wrap_client_query(self, client):
+        """Wrap ``ClaudeSDKClient.query``: not a run, but where a run's prompt goes out.
+
+        The interactive client sends a prompt with ``query()`` and reads the
+        reply with ``receive_response()``. The prompt is queued on the client
+        here and taken, oldest first, as a response's input there; a streamed
+        prompt is teed the same way as on the one-shot path.
+        """
+        query = getattr(client, "query", None)
+        if query is None or getattr(query, self._MARKER, False):
+            return
+
+        @wrapt.decorator
+        async def query_wrapper(wrapped, instance, args, kwargs):
+            prompt = kwargs.get("prompt", args[0] if args else None)
+            tee = _content.PromptTee.wrap(prompt)
+            if tee is not None:
+                args, kwargs = _content.replace_prompt(args, kwargs, tee)
+            result = await wrapped(*args, **kwargs)
+            # Queued only once the SDK accepted it: a prompt that never went
+            # out is not part of any response's input.
+            _content.queue_prompt(
+                instance,
+                tee.result() if tee is not None else _content.prompt_messages(prompt),
+            )
+            return result
+
+        wrapped = query_wrapper(query)
+        setattr(wrapped, self._MARKER, True)
+        client.query = wrapped
+
+    def _wrap_client(self, module):
+        """Wrap ``ClaudeSDKClient.receive_response`` (the interactive chokepoint)."""
+        client = getattr(module, "ClaudeSDKClient", None)
+        if client is None:
+            return
+        self._wrap_client_query(client)
+        receive_response = getattr(client, "receive_response", None)
+        if receive_response is None or getattr(receive_response, self._MARKER, False):
+            return
+
+        telemetry = self.telemetry
+
+        @wrapt.decorator
+        async def receive_response_wrapper(wrapped, instance, args, kwargs):
+            # The client keeps its ClaudeAgentOptions on the instance.
+            options = getattr(instance, "options", None)
+            details = _options_details(options)
+
+            with trace_agent_run(
+                telemetry,
+                semconv.span_name(semconv.OP_INVOKE_AGENT, "Claude Agent"),
+                {
+                    semconv.GEN_AI_OPERATION_NAME: semconv.OP_INVOKE_AGENT,
+                    semconv.GEN_AI_AGENT_NAME: "Claude Agent",
+                    semconv.GEN_AI_REQUEST_MODEL: details["model"],
+                    semconv.GEN_AI_TOOL_DEFINITIONS: semconv.tool_definitions(
+                        details["allowed_tools"]
+                    ),
+                },
+                start_event="agent_start",
+                end_event="agent_end",
+                error_event="agent_error",
+                log_fields={
+                    "agent_name": "Claude Agent",
+                    "agent_role": "Claude Agent",
+                    "model": details["model"],
+                },
+                start_details=details,
+            ) as run:
+                # The prompt went out earlier through ``client.query()``, which
+                # queued it on the client; the oldest prompt still unanswered
+                # is this response's input (one turn per prompt, see
+                # ``_content.take_prompt``).
+                run.set_input(_content.take_prompt(instance))
+                acc = UsageAccumulator()
+                out = _content.OutputAccumulator()
+                try:
+                    async for message in wrapped(*args, **kwargs):
+                        acc.add(extract_usage(getattr(message, "usage", None)))
+                        out.add(message)
+                        yield message
+                finally:
+                    run.set_usage(acc.result())
+                    run.set_output(out.result(), finish_reasons=out.finish_reasons)
+
+        wrapped = receive_response_wrapper(receive_response)
+        setattr(wrapped, self._MARKER, True)
+        client.receive_response = wrapped
+
+
+def install_claude_agent_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 Claude Agent SDK monitoring instrumentation.
+
+    Instruments the SDK's two async-generator entry points --
+    ``claude_agent_sdk.query(...)`` (one-shot) and
+    ``ClaudeSDKClient.receive_response()`` (interactive) -- so every agent run
+    streams start/end/error telemetry with zero changes to the agent code. Call
+    this *before* ``claude_agent_sdk`` is imported.
+
+    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 derived from the query prompt)
+    """
+    return install_hook(
+        ClaudeAgentSDKImportHook,
+        framework="Claude Agent SDK",
+        label="Claude Agent SDK",
+        server_url=server_url,
+        user=user,
+        api_key=api_key,
+        session=session,
+        name=name,
+    )

otelzero/hooks/crewai.py

@@ -0,0 +1,338 @@
+"""CrewAI instrumentation: wrap Crew/Agent/Task entry points."""
+
+from typing import Any, Optional
+
+import wrapt
+
+from otelzero import semconv
+from otelzero.hooks import _content
+from otelzero.hooks._usage import extract_usage
+from otelzero.hooks.base import _ImportInstrumenter, install_hook, trace_agent_run
+
+
+def _crewai_model_id(agent: Any) -> Optional[str]:
+    """Best-effort model identifier for a CrewAI agent's ``llm`` field.
+
+    ``Agent.llm`` is either a bare model string (the modern idiom, e.g.
+    ``"gpt-4o-mini"``) or a CrewAI ``LLM`` object whose ``.model`` attribute
+    holds the name. Returns ``None`` when unset -- the agent then falls back to
+    CrewAI's default LLM, which has no id we can read here.
+    """
+    llm = getattr(agent, "llm", None)
+    if llm is None:
+        return None
+    if isinstance(llm, str):
+        return llm or None
+    model = getattr(llm, "model", None)
+    if isinstance(model, str) and model:
+        return model
+    return None
+
+
+def _crew_model_id(instance: Any) -> Optional[str]:
+    """The shared model of a crew's agents, if they all agree.
+
+    A crew can mix models across agents; a single crew-level ``model`` is only
+    meaningful when every agent resolves to the same one, so return that when the
+    distinct set is a singleton and ``None`` otherwise.
+    """
+    models = {
+        m for a in (getattr(instance, "agents", None) or []) if (m := _crewai_model_id(a))
+    }
+    return next(iter(models)) if len(models) == 1 else None
+
+
+def _executor_messages(agent: Any) -> Optional[list]:
+    """The model messages CrewAI's executor accumulated for this agent.
+
+    ``Agent.agent_executor`` is built when a task is executed and keeps the
+    exchange in OpenAI's shape: a system turn, the task prompt as ``user``, an
+    assistant turn whose ``tool_calls`` name the tools it wants, each result
+    under the ``tool`` role, then the final assistant turn.
+    """
+    messages = getattr(getattr(agent, "agent_executor", None), "messages", None)
+    return messages if isinstance(messages, list) else None
+
+
+def _role_of(message: Any) -> Optional[str]:
+    if isinstance(message, dict):
+        return message.get("role")
+    return getattr(message, "role", None)
+
+
+def _content_of(message: Any) -> Any:
+    if isinstance(message, dict):
+        return message.get("content")
+    return getattr(message, "content", None)
+
+
+def _split_run(history: Optional[list], before: int):
+    """Split an executor's messages into system prompt, input and output.
+
+    The executor keeps the whole exchange, so a run that called a tool can be
+    recorded with its ``tool_call`` and ``tool_call_response`` parts rather than
+    as the answer text alone -- the agent's final string says nothing about the
+    tools that produced it. The system turn becomes the system instructions,
+    everything up to and including the task prompt is the input, and the
+    assistant turns with the tool results they triggered are the output.
+
+    ``None`` when there is nothing to split, so the caller keeps its fallback.
+    """
+    if not isinstance(history, list) or not 0 <= before < len(history):
+        return None
+    prior = [m for m in history[:before] if _role_of(m) != "system"]
+    new = history[before:]
+    system = [m for m in new if _role_of(m) == "system"]
+    rest = [m for m in new if _role_of(m) != "system"]
+    lead = 0
+    while lead < len(rest) and _role_of(rest[lead]) == "user":
+        lead += 1
+    return system, prior + rest[:lead], rest[lead:]
+
+
+class CrewAIImportHook(_ImportInstrumenter):
+    """Intercept CrewAI imports and instrument Crew/Agent/Task."""
+
+    WATCHED = {
+        "crewai.crew": "_wrap_crew_class",
+        "crewai.agent": "_wrap_agent_class",
+        "crewai.task": "_wrap_task_class",
+    }
+
+    def _wrap_crew_class(self, module):
+        """Wrap Crew class methods"""
+        if not hasattr(module, "Crew"):
+            return
+
+        original_kickoff = module.Crew.kickoff
+        telemetry = self.telemetry
+
+        @wrapt.decorator
+        def kickoff_wrapper(wrapped, instance, args, kwargs):
+            crew_model = _crew_model_id(instance)
+            crew_info = {
+                "num_agents": len(instance.agents)
+                if hasattr(instance, "agents")
+                else 0,
+                "num_tasks": len(instance.tasks) if hasattr(instance, "tasks") else 0,
+                "model": crew_model,
+            }
+            # Name the session after the crew's first task -- the run's headline
+            # intent -- so it reads e.g. "Research competitor pricing".
+            tasks = getattr(instance, "tasks", None) or []
+            crew_hint = getattr(tasks[0], "description", None) if tasks else None
+            # A crew coordinates several agents, which is exactly what the
+            # conventions call a workflow -- they name CrewAI as the example.
+            crew_name = getattr(instance, "name", None)
+
+            with trace_agent_run(
+                telemetry,
+                semconv.span_name(semconv.OP_INVOKE_WORKFLOW, crew_name),
+                {
+                    semconv.GEN_AI_OPERATION_NAME: semconv.OP_INVOKE_WORKFLOW,
+                    semconv.GEN_AI_WORKFLOW_NAME: crew_name,
+                    semconv.GEN_AI_REQUEST_MODEL: crew_model,
+                    semconv.OTELZERO_NUM_AGENTS: crew_info["num_agents"],
+                    semconv.OTELZERO_NUM_TASKS: crew_info["num_tasks"],
+                },
+                start_event="crew_start",
+                end_event="crew_end",
+                error_event="crew_error",
+                log_fields={"model": crew_model},
+                start_details=crew_info,
+                end_details={"status": "completed"},
+                session_name_hint=crew_hint,
+            ) as run:
+                # ``kickoff(inputs={...})`` -- the template variables are the
+                # crew's prompt; recorded as their JSON.
+                inputs = kwargs.get("inputs", args[0] if args else None)
+                run.set_input(_content.messages_from(inputs))
+                result = wrapped(*args, **kwargs)
+                # CrewOutput.token_usage is a UsageMetrics (prompt_tokens,
+                # completion_tokens, cached_prompt_tokens, total_tokens);
+                # CrewOutput.raw is the final answer.
+                run.set_usage(extract_usage(getattr(result, "token_usage", None)))
+                run.set_output(_content.text_of(result))
+                return result
+
+        module.Crew.kickoff = kickoff_wrapper(original_kickoff)
+
+    def _wrap_agent_class(self, module):
+        """Wrap Agent class methods"""
+        if not hasattr(module, "Agent"):
+            return
+
+        # Wrap execute_task if it exists
+        if hasattr(module.Agent, "execute_task"):
+            original_execute = module.Agent.execute_task
+            telemetry = self.telemetry
+
+            @wrapt.decorator
+            def execute_wrapper(wrapped, instance, args, kwargs):
+                agent_info = {
+                    "role": getattr(instance, "role", "Unknown"),
+                    "goal": getattr(instance, "goal", "N/A")[:100],
+                    "model": _crewai_model_id(instance),
+                }
+
+                with trace_agent_run(
+                    telemetry,
+                    semconv.span_name(semconv.OP_INVOKE_AGENT, agent_info["role"]),
+                    {
+                        semconv.GEN_AI_OPERATION_NAME: semconv.OP_INVOKE_AGENT,
+                        semconv.GEN_AI_AGENT_NAME: agent_info["role"],
+                        semconv.GEN_AI_REQUEST_MODEL: agent_info["model"],
+                        semconv.GEN_AI_AGENT_DESCRIPTION: agent_info["goal"],
+                    },
+                    start_event="agent_start",
+                    end_event="agent_end",
+                    error_event="agent_error",
+                    log_fields={
+                        "agent_role": agent_info["role"],
+                        "model": agent_info["model"],
+                    },
+                    start_details=agent_info,
+                    # Raw goal (not the "N/A"-defaulted copy above) names a
+                    # standalone agent run; a crew run is already named by its
+                    # first task, and the first non-empty hint wins.
+                    session_name_hint=getattr(instance, "goal", None),
+                ) as run:
+                    # ``execute_task(task, context=None, tools=None)``: the task
+                    # description (plus any context from earlier tasks) is the
+                    # prompt; the agent's backstory is its system prompt.
+                    task = kwargs.get("task", args[0] if args else None)
+                    context = kwargs.get("context", args[1] if len(args) > 1 else None)
+                    prompt = getattr(task, "description", None)
+                    if isinstance(prompt, str) and isinstance(context, str) and context:
+                        prompt = f"{prompt}\n\nContext:\n{context}"
+                    run.set_input(_content.messages_from(prompt))
+                    run.set_system(_content.text_of(getattr(instance, "backstory", None)))
+                    # Watermark the executor's history so this run's turns can
+                    # be told from the ones it inherited; a fresh executor for
+                    # this task starts the count over.
+                    history = _executor_messages(instance)
+                    before = len(history) if history is not None else 0
+                    result = wrapped(*args, **kwargs)
+                    after = _executor_messages(instance)
+                    split = _split_run(after, before if after is history else 0)
+                    messages = None
+                    if split is not None:
+                        system, inputs, outputs = split
+                        if system:
+                            # The prompt actually sent, which carries the role
+                            # and goal the backstory alone does not.
+                            run.set_system(_content.text_of(_content_of(system[0])))
+                        run.set_input(_content.messages_from(inputs))
+                        messages = _content.messages_from(
+                            outputs, default_role="assistant"
+                        )
+                    run.set_output(messages or _content.text_of(result))
+                    return result
+
+            module.Agent.execute_task = execute_wrapper(original_execute)
+
+    def _wrap_task_class(self, module):
+        """Wrap Task class methods"""
+        if not hasattr(module, "Task"):
+            return
+
+        telemetry = self.telemetry
+
+        def _traced(instance, args, kwargs):
+            """Build the trace context manager for one task execution."""
+            task_info = {
+                "description": getattr(instance, "description", "N/A")[:100],
+                "expected_output": getattr(instance, "expected_output", "N/A")[:100],
+            }
+
+            # The agent passed to this execution wins over the task's own, as it
+            # does inside CrewAI (a hierarchical crew assigns one at run time).
+            agent = (args[0] if args else kwargs.get("agent")) or getattr(
+                instance, "agent", None
+            )
+            agent_role = None
+            model = None
+            if agent:
+                agent_role = getattr(agent, "role", "Unknown")
+                model = _crewai_model_id(agent)
+
+            # A CrewAI Task is the *parent* of the agent invocation it
+            # delegates to, so no predefined operation fits; execute_task is
+            # otelzero's documented custom value (see otelzero.semconv).
+            return trace_agent_run(
+                telemetry,
+                semconv.span_name(semconv.OP_EXECUTE_TASK, agent_role),
+                {
+                    semconv.GEN_AI_OPERATION_NAME: semconv.OP_EXECUTE_TASK,
+                    semconv.GEN_AI_AGENT_NAME: agent_role,
+                    semconv.GEN_AI_REQUEST_MODEL: model,
+                    semconv.OTELZERO_TASK_DESCRIPTION: task_info["description"],
+                    semconv.OTELZERO_TASK_EXPECTED_OUTPUT: task_info["expected_output"],
+                },
+                start_event="task_start",
+                end_event="task_end",
+                error_event="task_error",
+                log_fields={"agent_role": agent_role, "model": model},
+                start_details=task_info,
+                end_details={"status": "completed"},
+                session_name_hint=getattr(instance, "description", None),
+            )
+
+        def _record_input(run, instance):
+            # The full description (the [:100] copy above is a preview).
+            run.set_input(_content.messages_from(getattr(instance, "description", None)))
+
+        @wrapt.decorator
+        def execute_wrapper(wrapped, instance, args, kwargs):
+            with _traced(instance, args, kwargs) as run:
+                _record_input(run, instance)
+                result = wrapped(*args, **kwargs)
+                # TaskOutput.raw is the task's result text.
+                run.set_output(_content.text_of(result))
+                return result
+
+        @wrapt.decorator
+        async def aexecute_wrapper(wrapped, instance, args, kwargs):
+            with _traced(instance, args, kwargs) as run:
+                _record_input(run, instance)
+                result = await wrapped(*args, **kwargs)
+                run.set_output(_content.text_of(result))
+                return result
+
+        # Current CrewAI has no Task.execute: execute_sync and execute_async
+        # both run _execute_core, and aexecute_sync runs _aexecute_core.
+        for name in ("_execute_core", "execute_sync", "execute"):
+            if hasattr(module.Task, name):
+                setattr(module.Task, name, execute_wrapper(getattr(module.Task, name)))
+                break
+        if hasattr(module.Task, "_aexecute_core"):
+            module.Task._aexecute_core = aexecute_wrapper(module.Task._aexecute_core)
+
+
+def install_crewai_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 CrewAI monitoring instrumentation
+
+    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 derived from the crew's first task)
+    """
+    return install_hook(
+        CrewAIImportHook,
+        framework="CrewAI",
+        label="CrewAI",
+        server_url=server_url,
+        user=user,
+        api_key=api_key,
+        session=session,
+        name=name,
+    )