agentloop-profiler 0.4.0__py3-none-any.whl

agentloop/integrations/openai.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import inspect
+import time
+from functools import wraps
+from typing import Any
+
+from agentloop.events import utc_now_iso
+from agentloop.tracer import current_event_id, current_trace, record_model_call
+
+# Marker set on wrappers so repeated instrumentation of the same callable is a
+# no-op instead of stacking wrappers and double-recording.
+_INSTRUMENTED_FLAG = "__agentloop_instrumented__"
+
+
+def _usage_get(usage: Any, *names: str) -> int:
+    for name in names:
+        if usage is None:
+            continue
+        if isinstance(usage, dict) and name in usage:
+            return int(usage.get(name) or 0)
+        value = getattr(usage, name, None)
+        if value is not None:
+            return int(value or 0)
+    return 0
+
+
+def _find_usage(obj: Any) -> Any:
+    """Locate a usage object on a response or a streaming event.
+
+    Non-streaming responses and chat chunks expose ``usage`` directly. The locked
+    Responses streaming API instead delivers final usage on the terminal
+    ``ResponseCompletedEvent`` under ``event.response.usage`` (the event itself has
+    no ``usage``), so fall back to the nested response shape.
+    """
+
+    if obj is None:
+        return None
+    usage = getattr(obj, "usage", None)
+    if usage is None and isinstance(obj, dict):
+        usage = obj.get("usage")
+    if usage is not None:
+        return usage
+    response = getattr(obj, "response", None)
+    if response is None and isinstance(obj, dict):
+        response = obj.get("response")
+    if response is None:
+        return None
+    nested = getattr(response, "usage", None)
+    if nested is None and isinstance(response, dict):
+        nested = response.get("usage")
+    return nested
+
+
+def _extract_usage(result: Any) -> tuple[int, int]:
+    usage = _find_usage(result)
+    input_tokens = _usage_get(usage, "input_tokens", "prompt_tokens")
+    output_tokens = _usage_get(usage, "output_tokens", "completion_tokens")
+    return input_tokens, output_tokens
+
+
+def _metadata_from_call(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+    metadata: dict[str, Any] = {"integration": "openai"}
+    if "tools" in kwargs and kwargs["tools"] is not None:
+        try:
+            metadata["tool_count"] = len(kwargs["tools"])
+        except TypeError:
+            metadata["tool_count"] = 1
+    if "messages" in kwargs and kwargs["messages"] is not None:
+        try:
+            metadata["message_count"] = len(kwargs["messages"])
+        except TypeError:
+            metadata["message_count"] = 1
+    if args:
+        metadata["positional_args"] = len(args)
+    return metadata
+
+
+def _model_from_call(kwargs: dict[str, Any], default_model: str | None = None) -> str | None:
+    return (
+        str(kwargs.get("model") or default_model)
+        if (kwargs.get("model") or default_model)
+        else None
+    )
+
+
+class _CallRecorder:
+    """Collects timing for one instrumented call and records it exactly once.
+
+    Recording is deferred so streaming calls can be finalized when the stream is
+    consumed rather than when the iterator is created. Trace ownership is captured
+    at invocation time: the event is recorded into the trace that was active when
+    the call was made, regardless of which context consumes the stream later. If no
+    trace was active at invocation the call is never recorded, and the application
+    call is not disturbed.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+        default_model: str | None,
+    ) -> None:
+        self._name = name
+        self._args = args
+        self._kwargs = kwargs
+        self._default_model = default_model
+        self._started_at = utc_now_iso()
+        self._start = time.perf_counter()
+        self._done = False
+        # Capture ownership now, at invocation time, not at finalization time.
+        self._trace = current_trace()
+        self._parent_id = current_event_id()
+
+    def finalize(self, result: Any = None, *, status: str = "ok", error: str | None = None) -> None:
+        if self._done:
+            return
+        self._done = True
+        if self._trace is None:
+            # No trace was active when the call was invoked: never record, and do
+            # not turn the caller's successful application call into an error.
+            return
+        input_tokens, output_tokens = _extract_usage(result)
+        record_model_call(
+            self._name,
+            started_at=self._started_at,
+            duration_ms=(time.perf_counter() - self._start) * 1000,
+            model=_model_from_call(self._kwargs, self._default_model),
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            status=status,
+            error=error,
+            metadata=_metadata_from_call(self._args, self._kwargs),
+            parent_id=self._parent_id,
+            trace=self._trace,
+        )
+
+
+def _has_usage(obj: Any) -> bool:
+    return _find_usage(obj) is not None
+
+
+def _stream_requested(kwargs: dict[str, Any]) -> bool:
+    return bool(kwargs.get("stream"))
+
+
+def _is_sync_stream(result: Any) -> bool:
+    if isinstance(result, (dict, list, tuple, set, str, bytes, bytearray)):
+        return False
+    return (
+        hasattr(result, "__iter__") or hasattr(result, "__next__") or hasattr(result, "__enter__")
+    )
+
+
+def _is_async_stream(result: Any) -> bool:
+    return (
+        hasattr(result, "__aiter__")
+        or hasattr(result, "__anext__")
+        or hasattr(result, "__aenter__")
+        or inspect.isasyncgen(result)
+    )
+
+
+class _InstrumentedSyncStream:
+    """Proxy over a sync streaming response that finalizes on consumption.
+
+    Iteration, context-manager use, and ``close()`` all end the stream; whichever
+    happens first records the call (once). Every other attribute is delegated to
+    the wrapped stream so SDK behavior is preserved.
+    """
+
+    def __init__(self, stream: Any, recorder: _CallRecorder) -> None:
+        self._stream = stream
+        self._recorder = recorder
+        self._iter: Any = None
+        self._usage_source: Any = None
+
+    def __iter__(self) -> _InstrumentedSyncStream:
+        self._iter = iter(self._stream)
+        return self
+
+    def __next__(self) -> Any:
+        if self._iter is None:
+            self._iter = iter(self._stream)
+        try:
+            chunk = next(self._iter)
+        except StopIteration:
+            self._recorder.finalize(self._usage_source)
+            raise
+        except BaseException as exc:  # noqa: BLE001 - record then propagate unchanged
+            self._recorder.finalize(self._usage_source, status="error", error=str(exc))
+            raise
+        if _has_usage(chunk):
+            self._usage_source = chunk
+        return chunk
+
+    def __enter__(self) -> _InstrumentedSyncStream:
+        if hasattr(self._stream, "__enter__"):
+            self._stream.__enter__()
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+        suppress = None
+        if hasattr(self._stream, "__exit__"):
+            suppress = self._stream.__exit__(exc_type, exc, tb)
+        if exc is not None:
+            self._recorder.finalize(self._usage_source, status="error", error=str(exc))
+        else:
+            self._recorder.finalize(self._usage_source)
+        return suppress
+
+    def close(self) -> Any:
+        try:
+            if hasattr(self._stream, "close"):
+                return self._stream.close()
+            return None
+        finally:
+            self._recorder.finalize(self._usage_source)
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._stream, item)
+
+
+class _InstrumentedAsyncStream:
+    """Async counterpart of :class:`_InstrumentedSyncStream`."""
+
+    def __init__(self, stream: Any, recorder: _CallRecorder) -> None:
+        self._stream = stream
+        self._recorder = recorder
+        self._iter: Any = None
+        self._usage_source: Any = None
+
+    def __aiter__(self) -> _InstrumentedAsyncStream:
+        self._iter = (
+            self._stream.__aiter__() if hasattr(self._stream, "__aiter__") else self._stream
+        )
+        return self
+
+    async def __anext__(self) -> Any:
+        if self._iter is None:
+            self._iter = (
+                self._stream.__aiter__() if hasattr(self._stream, "__aiter__") else self._stream
+            )
+        try:
+            chunk = await self._iter.__anext__()
+        except StopAsyncIteration:
+            self._recorder.finalize(self._usage_source)
+            raise
+        except BaseException as exc:  # noqa: BLE001 - record then propagate unchanged
+            self._recorder.finalize(self._usage_source, status="error", error=str(exc))
+            raise
+        if _has_usage(chunk):
+            self._usage_source = chunk
+        return chunk
+
+    async def __aenter__(self) -> _InstrumentedAsyncStream:
+        if hasattr(self._stream, "__aenter__"):
+            await self._stream.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+        suppress = None
+        if hasattr(self._stream, "__aexit__"):
+            suppress = await self._stream.__aexit__(exc_type, exc, tb)
+        if exc is not None:
+            self._recorder.finalize(self._usage_source, status="error", error=str(exc))
+        else:
+            self._recorder.finalize(self._usage_source)
+        return suppress
+
+    async def aclose(self) -> Any:
+        try:
+            if hasattr(self._stream, "aclose"):
+                return await self._stream.aclose()
+            return None
+        finally:
+            self._recorder.finalize(self._usage_source)
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._stream, item)
+
+
+def instrument_callable(
+    fn: Any,
+    *,
+    name: str = "openai.model_call",
+    default_model: str | None = None,
+) -> Any:
+    """Wrap a sync or async OpenAI-like callable and record token usage.
+
+    The wrapper is dependency-free and works with OpenAI SDK resource methods such
+    as `client.responses.create(...)` or `client.chat.completions.create(...)`.
+    It records prompt/completion usage when the returned object exposes a
+    `.usage` object or a `{"usage": ...}` dict.
+
+    Wrapping is idempotent: re-instrumenting an already-instrumented callable
+    returns it unchanged, so one request records one event. Streaming responses
+    (requested with `stream=True`) are finalized when the stream is consumed,
+    closed, fails, or is cancelled — not when the iterator is created — so the
+    recorded duration covers consumption and final usage is captured when the SDK
+    provides it. Calls made without an active trace are not recorded and are not
+    turned into errors.
+    """
+
+    if getattr(fn, _INSTRUMENTED_FLAG, False):
+        return fn
+
+    if inspect.iscoroutinefunction(fn):
+
+        @wraps(fn)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            recorder = _CallRecorder(name, args, kwargs, default_model)
+            try:
+                result = await fn(*args, **kwargs)
+            except Exception as exc:
+                recorder.finalize(None, status="error", error=str(exc))
+                raise
+            if _stream_requested(kwargs) and _is_async_stream(result):
+                return _InstrumentedAsyncStream(result, recorder)
+            recorder.finalize(result)
+            return result
+
+        setattr(async_wrapper, _INSTRUMENTED_FLAG, True)
+        return async_wrapper
+
+    @wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        recorder = _CallRecorder(name, args, kwargs, default_model)
+        try:
+            result = fn(*args, **kwargs)
+        except Exception as exc:
+            recorder.finalize(None, status="error", error=str(exc))
+            raise
+        if _stream_requested(kwargs) and _is_sync_stream(result):
+            return _InstrumentedSyncStream(result, recorder)
+        recorder.finalize(result)
+        return result
+
+    setattr(wrapper, _INSTRUMENTED_FLAG, True)
+    return wrapper
+
+
+def instrument_openai_client(client: Any) -> Any:
+    """Patch common OpenAI SDK call sites in-place.
+
+    Supported call sites when present:
+    - client.responses.create
+    - client.chat.completions.create
+
+    The function returns the same client object for ergonomic usage:
+
+    ```python
+    client = instrument_openai_client(OpenAI())
+    ```
+    """
+
+    responses = getattr(client, "responses", None)
+    if responses is not None and callable(getattr(responses, "create", None)):
+        responses.create = instrument_callable(responses.create, name="openai.responses.create")
+
+    chat = getattr(client, "chat", None)
+    completions = getattr(chat, "completions", None) if chat is not None else None
+    if completions is not None and callable(getattr(completions, "create", None)):
+        completions.create = instrument_callable(
+            completions.create, name="openai.chat.completions.create"
+        )
+
+    return client

agentloop/integrations/openai_agents.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from agentloop.otel import trace_from_otel
+from agentloop.otel_ids import to_span_id, to_trace_id
+from agentloop.tracer import AgentTrace
+
+
+class AgentLoopTracingProcessor:
+    """OpenAI Agents SDK-compatible tracing processor.
+
+    The OpenAI Agents SDK calls `on_trace_start`, `on_trace_end`, `on_span_start`,
+    and `on_span_end` on registered processors. This class avoids importing the SDK
+    so AgentLoop can expose the integration without adding a hard dependency.
+    """
+
+    def __init__(self, out_dir: str | Path | None = None, *, retain_exported: bool = True):
+        self.out_dir = Path(out_dir) if out_dir is not None else None
+        # Retain exported AgentTrace objects on ``exported_traces``. Long-lived
+        # processors that never read that list can set ``retain_exported=False``
+        # so completed traces are not held in memory.
+        self.retain_exported = retain_exported
+        self._traces: dict[str, Any] = {}
+        # Trace ids started but not yet ended, in start order. Used to attribute
+        # spans that arrive without a readable trace id to a single open trace.
+        self._active_trace_ids: list[str] = []
+        # Spans are grouped by their owning trace id and released when that trace
+        # ends, so memory does not grow with the number of completed traces.
+        self._spans_by_trace: dict[str, list[Any]] = {}
+        self.exported_traces: list[AgentTrace] = []
+
+    def on_trace_start(self, trace: Any) -> None:
+        trace_id = str(_read(trace, "trace_id", "id") or "trace_unknown")
+        self._traces[trace_id] = trace
+        if trace_id not in self._active_trace_ids:
+            self._active_trace_ids.append(trace_id)
+
+    def on_trace_end(self, trace: Any) -> None:
+        trace_id = str(_read(trace, "trace_id", "id") or "trace_unknown")
+        if not self._is_known(trace_id):
+            # Unknown trace id: already exported (duplicate on_trace_end) or never
+            # observed. Nothing to export and nothing to release.
+            return
+
+        try:
+            agent_trace = self._build_trace(trace, trace_id)
+            if self.retain_exported:
+                self.exported_traces.append(agent_trace)
+            if self.out_dir is not None:
+                self.out_dir.mkdir(parents=True, exist_ok=True)
+                agent_trace.export_json(self.out_dir / f"{agent_trace.run_id}.json")
+        finally:
+            # Release state for the completed trace even if build/export fails, so a
+            # failed export cannot re-introduce the completed-trace memory leak.
+            self._traces.pop(trace_id, None)
+            self._spans_by_trace.pop(trace_id, None)
+            if trace_id in self._active_trace_ids:
+                self._active_trace_ids.remove(trace_id)
+
+    def on_span_start(self, span: Any) -> None:
+        return None
+
+    def on_span_end(self, span: Any) -> None:
+        trace_id = _span_trace_id(span)
+        if trace_id is None:
+            # No readable trace id: attribute the span to the most recently started
+            # open trace. If none is open, drop it rather than copy it into an
+            # unrelated trace that ends later.
+            trace_id = self._current_active_trace_id()
+            if trace_id is None:
+                return
+        self._spans_by_trace.setdefault(trace_id, []).append(span)
+
+    def shutdown(self) -> None:
+        # Release all retained state; further callbacks start from empty.
+        self._traces.clear()
+        self._active_trace_ids.clear()
+        self._spans_by_trace.clear()
+
+    def force_flush(self) -> None:
+        # Spans are exported when their trace ends; nothing is buffered to flush.
+        return None
+
+    def _current_active_trace_id(self) -> str | None:
+        return self._active_trace_ids[-1] if self._active_trace_ids else None
+
+    def _is_known(self, trace_id: str) -> bool:
+        return (
+            trace_id in self._traces
+            or trace_id in self._spans_by_trace
+            or trace_id in self._active_trace_ids
+        )
+
+    def _build_trace(self, trace: Any, trace_id: str | None = None) -> AgentTrace:
+        if trace_id is None:
+            trace_id = str(_read(trace, "trace_id", "id") or "trace_unknown")
+        name = str(_read(trace, "name", "workflow_name") or "openai_agents_trace")
+        spans = [_span_to_otel(span, trace_id) for span in self._spans_by_trace.get(trace_id, [])]
+        return trace_from_otel({"spans": spans}, name=name)
+
+
+def _span_to_otel(span: Any, trace_id: str) -> dict[str, Any]:
+    span_data = _read(span, "span_data") or _read(span, "data") or {}
+    span_type = _read(span_data, "type") or span_data.__class__.__name__.lower()
+    name = str(_read(span_data, "name") or _read(span, "name") or span_type or "openai_agents_span")
+    model = _read(span_data, "model")
+    usage = _read(span_data, "usage") or {}
+    attrs = [
+        _attribute("gen_ai.operation.name", _operation_from_span_type(str(span_type))),
+        _attribute("agentloop.event_type", _event_type_from_span_type(str(span_type))),
+        _attribute("agentloop.name", name),
+    ]
+    if model:
+        attrs.append(_attribute("gen_ai.request.model", str(model)))
+    input_tokens = (
+        _read(usage, "input_tokens", "prompt_tokens")
+        or _read(span_data, "input_tokens", "prompt_tokens")
+        or 0
+    )
+    output_tokens = (
+        _read(usage, "output_tokens", "completion_tokens")
+        or _read(span_data, "output_tokens", "completion_tokens")
+        or 0
+    )
+    attrs.append(_attribute("gen_ai.usage.input_tokens", int(input_tokens or 0)))
+    attrs.append(_attribute("gen_ai.usage.output_tokens", int(output_tokens or 0)))
+
+    native_span_id = str(_read(span, "span_id", "id") or "0")
+    # Preserve the original native ids for diagnosability (see agentloop.otel_ids).
+    attrs.append(_attribute("agentloop.native_span_id", native_span_id))
+    attrs.append(_attribute("agentloop.native_trace_id", str(trace_id)))
+    out = {
+        "traceId": to_trace_id(trace_id),
+        "spanId": to_span_id(native_span_id),
+        "name": name,
+        "startTimeUnixNano": str(_time_ns(_read(span, "started_at", "start_time"))),
+        "endTimeUnixNano": str(_time_ns(_read(span, "ended_at", "end_time"))),
+        "attributes": attrs,
+        "status": {"code": "STATUS_CODE_OK"},
+    }
+    parent_id = _read(span, "parent_id", "parent_span_id")
+    if parent_id:
+        attrs.append(_attribute("agentloop.native_parent_id", str(parent_id)))
+        out["parentSpanId"] = to_span_id(str(parent_id))
+    return out
+
+
+def _read(obj: Any, *names: str) -> Any:
+    for name in names:
+        if isinstance(obj, dict) and name in obj:
+            return obj[name]
+        if hasattr(obj, name):
+            return getattr(obj, name)
+    return None
+
+
+def _span_trace_id(span: Any) -> str | None:
+    value = _read(span, "trace_id")
+    return str(value) if value else None
+
+
+def _operation_from_span_type(span_type: str) -> str:
+    lower = span_type.lower()
+    if "generation" in lower or "response" in lower or "llm" in lower:
+        return "chat"
+    if "function" in lower or "tool" in lower:
+        return "execute_tool"
+    return "invoke_agent"
+
+
+def _event_type_from_span_type(span_type: str) -> str:
+    operation = _operation_from_span_type(span_type)
+    return "model_call" if operation == "chat" else "tool_call"
+
+
+def _attribute(key: str, value: Any) -> dict[str, Any]:
+    if isinstance(value, bool):
+        encoded = {"boolValue": value}
+    elif isinstance(value, int):
+        encoded = {"intValue": value}
+    elif isinstance(value, float):
+        encoded = {"doubleValue": value}
+    else:
+        encoded = {"stringValue": str(value)}
+    return {"key": key, "value": encoded}
+
+
+def _time_ns(value: Any) -> int:
+    if value is None:
+        return 0
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        return int(value * 1_000_000_000)
+    try:
+        from datetime import datetime
+
+        return int(
+            datetime.fromisoformat(str(value).replace("Z", "+00:00")).timestamp() * 1_000_000_000
+        )
+    except ValueError:
+        return 0
+
+
+def export_trace_json(trace: AgentTrace, path: str | Path) -> Path:
+    out = Path(path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps(trace.to_dict(), indent=2), encoding="utf-8")
+    return out

agentloop/integrations/sdk.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def instrument(obj: Any, *, kind: str | None = None, name: str | None = None) -> Any:
+    """Best-effort one-line instrumentation for common agent objects.
+
+    Explicit `kind` values:
+    - `openai`
+    - `crewai_crew`
+    - `crewai_agent`
+    - `crewai_task`
+    - `langgraph_state_graph`
+    - `langgraph_runnable`
+
+    Without `kind`, AgentLoop uses duck-typing to choose the safest wrapper.
+    """
+
+    normalized = kind.lower().replace("-", "_") if kind else None
+
+    if normalized == "openai" or _looks_like_openai_client(obj):
+        from agentloop.integrations.openai import instrument_openai_client
+
+        return instrument_openai_client(obj)
+
+    if normalized == "crewai_crew" or _looks_like_crewai_crew(obj):
+        from agentloop.integrations.crewai import instrument_crew
+
+        return instrument_crew(obj, name=name)
+
+    if normalized == "crewai_agent":
+        from agentloop.integrations.crewai import instrument_agent
+
+        return instrument_agent(obj, name=name)
+
+    if normalized == "crewai_task":
+        from agentloop.integrations.crewai import instrument_task
+
+        return instrument_task(obj, name=name)
+
+    if normalized == "langgraph_state_graph" or _looks_like_state_graph(obj):
+        from agentloop.integrations.langgraph import instrument_state_graph
+
+        return instrument_state_graph(obj)
+
+    if normalized == "langgraph_runnable" or _looks_like_runnable(obj):
+        from agentloop.integrations.langgraph import trace_runnable
+
+        return trace_runnable(
+            obj, name=name or getattr(obj, "name", None) or obj.__class__.__name__
+        )
+
+    raise ValueError(
+        "Could not infer object type for AgentLoop instrumentation. "
+        "Pass kind='openai', 'crewai_crew', 'crewai_agent', 'crewai_task', "
+        "'langgraph_state_graph', or 'langgraph_runnable'."
+    )
+
+
+def _looks_like_openai_client(obj: Any) -> bool:
+    responses = getattr(obj, "responses", None)
+    chat = getattr(obj, "chat", None)
+    completions = getattr(chat, "completions", None) if chat is not None else None
+    return bool(
+        responses is not None
+        and callable(getattr(responses, "create", None))
+        or completions is not None
+        and callable(getattr(completions, "create", None))
+    )
+
+
+def _looks_like_crewai_crew(obj: Any) -> bool:
+    return callable(getattr(obj, "kickoff", None)) or callable(getattr(obj, "kickoff_async", None))
+
+
+def _looks_like_state_graph(obj: Any) -> bool:
+    return callable(getattr(obj, "add_node", None)) and callable(getattr(obj, "compile", None))
+
+
+def _looks_like_runnable(obj: Any) -> bool:
+    return callable(getattr(obj, "invoke", None)) or callable(getattr(obj, "ainvoke", None))