brooder 0.1.0__py3-none-any.whl

brooder/integrations/langchain.py

@@ -0,0 +1,321 @@
+"""LangChain / LangGraph capture — trajectories from a callback handler.
+
+Attach Brooder's callback handler to a LangChain or LangGraph run and it records the agent's
+trajectory automatically — no manual ``tool_call``, and no OpenTelemetry instrumentation to wire:
+
+    import brooder.integrations.langchain as sg_lc
+
+    handler = sg_lc.callback_handler(agent="support-agent")
+    graph.invoke({"messages": [...]}, config={"callbacks": [handler]})
+
+Because LangGraph runs on LangChain's callback system, one handler covers both.
+
+**How callbacks map to steps** (verified against ``langchain_core`` 0.3.x):
+
+- the **root** ``on_chain_start`` (``parent_run_id is None``) opens a run; its input is the case id.
+- ``on_chat_model_start`` / ``on_llm_start`` → a ``TURN`` step.
+- ``on_tool_start`` + ``on_tool_end`` → a ``TOOL`` step (paired by the tool's own ``run_id``).
+- the root ``on_chain_end`` → a ``FINAL`` step and save.
+
+Nested LLM/tool calls carry a ``parent_run_id`` chain, so the handler maps every ``run_id`` to its
+root to attach steps to the right run. A sync handler still fires during async runs but may run on
+worker threads, so capture goes through the thread-safe :class:`~brooder.recorder.RunHandle`.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from typing import Any, Optional
+
+from .. import recorder
+from ..errors import BrooderError
+from ..log import get_logger
+from .base import as_dict, parse_json
+
+_log = get_logger()
+
+
+def _coerce(value: Any) -> Any:
+    """Coerce a value into a JSON-safe form (unknown objects become their ``str``)."""
+    try:
+        return json.loads(json.dumps(value, default=str))
+    except (TypeError, ValueError):
+        return str(value)
+
+
+def _message_content(message: Any) -> Any:
+    """Return a LangChain message's ``.content`` (or its ``str`` if it isn't a message)."""
+    content = getattr(message, "content", None)
+    return content if content is not None else str(message)
+
+
+def _extract_input(inputs: Any) -> Any:
+    """Recover a run's inputs (for case identity) from the root chain's ``inputs``."""
+    if isinstance(inputs, dict):
+        messages = inputs.get("messages")
+        if isinstance(messages, list) and messages:
+            return _coerce([_message_content(m) for m in messages])
+        if len(inputs) == 1:
+            return _coerce(next(iter(inputs.values())))
+    return _coerce(inputs)
+
+
+def _final_output(outputs: Any) -> Any:
+    """Recover the final answer from the root chain's ``outputs`` (LangGraph state / result)."""
+    if isinstance(outputs, dict):
+        messages = outputs.get("messages")
+        if isinstance(messages, list) and messages:
+            return _coerce(_message_content(messages[-1]))
+        if len(outputs) == 1:
+            return _coerce(next(iter(outputs.values())))
+    return _coerce(outputs)
+
+
+def _tool_name(serialized: Any, extra: dict[str, Any]) -> str:
+    """Resolve a tool's name from its ``serialized`` payload, falling back to callback kwargs."""
+    name = (serialized or {}).get("name") if isinstance(serialized, dict) else None
+    return name or extra.get("name") or "tool"
+
+
+def _tool_args(inputs: Any, input_str: Any) -> dict[str, Any]:
+    """Resolve a tool's arguments — the structured ``inputs`` dict if given, else ``input_str``."""
+    if inputs is not None:
+        return as_dict(inputs)
+    return parse_json(input_str)
+
+
+def _tool_output(output: Any) -> Any:
+    """Extract a tool result: a ``ToolMessage`` ``.content`` (0.2+) or the raw value, JSON-safe."""
+    content = getattr(output, "content", None)
+    return _coerce(content if content is not None else output)
+
+
+class _Capture:
+    """The framework-agnostic capture logic behind the callback handler.
+
+    Kept separate from the ``BaseCallbackHandler`` subclass so it imports and tests without any
+    LangChain install. Keys one run per root ``run_id``; maps each nested run id to its root; pairs
+    tool start/end by the tool's own ``run_id``. Thread-safe via the recorder's ``RunHandle`` and an
+    internal lock over its bookkeeping dicts.
+    """
+
+    def __init__(self, agent: Optional[str]) -> None:
+        self._agent = agent
+        self._root_of: dict[str, str] = {}
+        self._pending: dict[str, tuple[str, str, dict[str, Any]]] = {}
+        self._lock = threading.Lock()
+
+    def _resolve_root(self, run_id: Any, parent_run_id: Any) -> str:
+        """Record ``run_id``'s place in the run tree and return its root run id."""
+        rid = str(run_id)
+        pid = str(parent_run_id) if parent_run_id is not None else None
+        with self._lock:
+            if pid is None:
+                self._root_of[rid] = rid
+                return rid
+            root = self._root_of.get(pid, pid)
+            self._root_of[rid] = root
+            return root
+
+    def _cleanup(self, root: str) -> None:
+        """Drop bookkeeping for a finished root run."""
+        with self._lock:
+            self._root_of = {k: v for k, v in self._root_of.items() if v != root}
+            self._pending = {k: v for k, v in self._pending.items() if v[0] != root}
+
+    def chain_start(self, serialized: Any, inputs: Any, run_id: Any, parent_run_id: Any) -> None:
+        """Open a run when the root chain starts."""
+        try:
+            root = self._resolve_root(run_id, parent_run_id)
+            if parent_run_id is None and recorder.get_run(root) is None:
+                recorder.open_run(
+                    self._agent or "langchain-agent",
+                    inputs=_extract_input(inputs),
+                    external_id=root,
+                )
+        except Exception:  # capture must never break the app
+            _log.debug("brooder: langchain chain_start capture failed", exc_info=True)
+
+    def model_start(self, payload: Any, run_id: Any, parent_run_id: Any) -> None:
+        """Append a TURN step for a model call (chat or completion)."""
+        try:
+            root = self._resolve_root(run_id, parent_run_id)
+            handle = recorder.get_run(root)
+            if handle is None and parent_run_id is None:
+                handle = recorder.open_run(
+                    self._agent or "langchain-agent",
+                    inputs=_extract_input(payload),
+                    external_id=root,
+                )
+            if handle is not None:
+                handle.turn({"provider": "langchain"})
+        except Exception:
+            _log.debug("brooder: langchain model_start capture failed", exc_info=True)
+
+    def tool_start(
+        self, serialized: Any, input_str: Any, run_id: Any, parent_run_id: Any, inputs: Any
+    ) -> None:
+        """Stash a tool call, to be emitted with its result at :meth:`tool_end`."""
+        try:
+            root = self._resolve_root(run_id, parent_run_id)
+            name = _tool_name(serialized, {})
+            args = _tool_args(inputs, input_str)
+            with self._lock:
+                self._pending[str(run_id)] = (root, name, args)
+        except Exception:
+            _log.debug("brooder: langchain tool_start capture failed", exc_info=True)
+
+    def tool_end(self, output: Any, run_id: Any, parent_run_id: Any) -> None:
+        """Emit the TOOL step for a completed tool call."""
+        self._emit_tool(run_id, _tool_output(output))
+
+    def tool_error(self, error: Any, run_id: Any, parent_run_id: Any) -> None:
+        """Emit the TOOL step for a failed tool call (result is the error)."""
+        self._emit_tool(run_id, f"error: {error}")
+
+    def _emit_tool(self, run_id: Any, result: Any) -> None:
+        try:
+            with self._lock:
+                entry = self._pending.pop(str(run_id), None)
+            if entry is None:
+                return
+            root, name, args = entry
+            handle = recorder.get_run(root)
+            if handle is not None:
+                handle.tool_call(name, args, result=result)
+        except Exception:
+            _log.debug("brooder: langchain tool_end capture failed", exc_info=True)
+
+    def chain_end(self, outputs: Any, run_id: Any, parent_run_id: Any) -> None:
+        """Finalize the run when the root chain ends."""
+        self._finalize(run_id, parent_run_id, _final_output(outputs))
+
+    def chain_error(self, error: Any, run_id: Any, parent_run_id: Any) -> None:
+        """Finalize the run when the root chain errors (no final answer)."""
+        self._finalize(run_id, parent_run_id, None)
+
+    def _finalize(self, run_id: Any, parent_run_id: Any, output: Any) -> None:
+        try:
+            if parent_run_id is not None:
+                return  # only the root chain delimits the run
+            root = str(run_id)
+            handle = recorder.get_run(root)
+            if handle is not None:
+                handle.finish(output)
+            self._cleanup(root)
+        except Exception:
+            _log.debug("brooder: langchain chain_end capture failed", exc_info=True)
+
+
+_handler_cls: Optional[type] = None
+
+
+def _build_handler_class() -> type:
+    """Build the ``BaseCallbackHandler`` subclass lazily (needs ``langchain_core``)."""
+    from langchain_core.callbacks import BaseCallbackHandler
+
+    class BrooderCallbackHandler(BaseCallbackHandler):  # type: ignore[misc]
+        """A LangChain callback handler that records the run into a Brooder trajectory."""
+
+        raise_error = False
+
+        def __init__(self, agent: Optional[str] = None) -> None:
+            """Create a handler recording into runs named ``agent``."""
+            super().__init__()
+            self._cap = _Capture(agent)
+
+        def on_chain_start(
+            self, serialized: Any, inputs: Any, *, run_id: Any, parent_run_id: Any = None, **kw: Any
+        ) -> None:
+            """Forward a chain start to the capture logic."""
+            self._cap.chain_start(serialized, inputs, run_id, parent_run_id)
+
+        def on_chain_end(
+            self, outputs: Any, *, run_id: Any, parent_run_id: Any = None, **kw: Any
+        ) -> None:
+            """Forward a chain end to the capture logic."""
+            self._cap.chain_end(outputs, run_id, parent_run_id)
+
+        def on_chain_error(
+            self, error: Any, *, run_id: Any, parent_run_id: Any = None, **kw: Any
+        ) -> None:
+            """Forward a chain error to the capture logic."""
+            self._cap.chain_error(error, run_id, parent_run_id)
+
+        def on_chat_model_start(
+            self,
+            serialized: Any,
+            messages: Any,
+            *,
+            run_id: Any,
+            parent_run_id: Any = None,
+            **kw: Any,
+        ) -> None:
+            """Forward a chat-model start (a model turn) to the capture logic."""
+            self._cap.model_start(messages, run_id, parent_run_id)
+
+        def on_llm_start(
+            self,
+            serialized: Any,
+            prompts: Any,
+            *,
+            run_id: Any,
+            parent_run_id: Any = None,
+            **kw: Any,
+        ) -> None:
+            """Forward a completion-model start (a model turn) to the capture logic."""
+            self._cap.model_start(prompts, run_id, parent_run_id)
+
+        def on_tool_start(
+            self,
+            serialized: Any,
+            input_str: Any,
+            *,
+            run_id: Any,
+            parent_run_id: Any = None,
+            inputs: Any = None,
+            **kw: Any,
+        ) -> None:
+            """Forward a tool start to the capture logic."""
+            self._cap.tool_start(serialized, input_str, run_id, parent_run_id, inputs)
+
+        def on_tool_end(
+            self, output: Any, *, run_id: Any, parent_run_id: Any = None, **kw: Any
+        ) -> None:
+            """Forward a tool end to the capture logic."""
+            self._cap.tool_end(output, run_id, parent_run_id)
+
+        def on_tool_error(
+            self, error: Any, *, run_id: Any, parent_run_id: Any = None, **kw: Any
+        ) -> None:
+            """Forward a tool error to the capture logic."""
+            self._cap.tool_error(error, run_id, parent_run_id)
+
+    return BrooderCallbackHandler
+
+
+def callback_handler(agent: Optional[str] = None) -> Any:
+    """Build a Brooder LangChain callback handler for ``config={"callbacks": [handler]}``.
+
+    Records the run's tool trajectory into a Brooder run (see the module docstring for details).
+    Covers LangGraph too, since it runs on LangChain's callbacks.
+
+    Args:
+        agent: Logical agent name used to group baselines (defaults to ``"langchain-agent"``).
+
+    Returns:
+        A ``langchain_core`` ``BaseCallbackHandler`` instance.
+
+    Raises:
+        BrooderError: If LangChain (``langchain_core``) is not installed.
+    """
+    global _handler_cls
+    if _handler_cls is None:
+        try:
+            _handler_cls = _build_handler_class()
+        except ImportError as exc:
+            raise BrooderError(
+                "langchain.callback_handler() needs LangChain — `pip install langchain-core`"
+            ) from exc
+    return _handler_cls(agent)

brooder/integrations/openai.py

@@ -0,0 +1,43 @@
+"""OpenAI and Azure OpenAI auto-capture.
+
+Both use the ``openai`` SDK, so one normalizer covers ``OpenAI``, ``AzureOpenAI``, and
+API-compatible clients. Wraps ``client.chat.completions.create``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from .base import NormalizedCall, ToolRequest, get, parse_json, wrap
+
+
+def _normalize(_kwargs: dict[str, Any], response: Any) -> Optional[NormalizedCall]:
+    choices = get(response, "choices") or []
+    message = get(choices[0], "message") if choices else None
+    tool_calls = [
+        ToolRequest(
+            name=get(get(tc, "function"), "name") or "",
+            arguments=parse_json(get(get(tc, "function"), "arguments")),
+        )
+        for tc in (get(message, "tool_calls") or [])
+    ]
+    return NormalizedCall(
+        provider="openai",
+        model=get(response, "model"),
+        content=get(message, "content"),
+        tool_calls=tool_calls,
+    )
+
+
+def instrument(client: Any, capture_content: bool = False) -> Any:
+    """Instrument an OpenAI/Azure OpenAI client's ``chat.completions.create``.
+
+    Args:
+        client: An ``OpenAI`` or ``AzureOpenAI`` client instance.
+        capture_content: Also record assistant text, not just tool calls.
+
+    Returns:
+        The same client, patched in place.
+    """
+    wrap(client.chat.completions, "create", _normalize, capture_content)
+    return client

brooder/integrations/openai_agents.py

@@ -0,0 +1,208 @@
+"""OpenAI Agents SDK capture — trajectories from the SDK's tracing.
+
+Its tracing is on by default. Install Brooder's trace processor and it records the agent's
+trajectory automatically — no manual ``tool_call``:
+
+    import brooder.integrations.openai_agents as sg_agents
+    sg_agents.install(agent="support-agent")   # then run your agent as usual
+
+    from agents import Agent, Runner
+    Runner.run_sync(Agent(name="support", ...), "cancel my order")
+
+`install` appends the processor (the default OpenAI exporter keeps working) and needs no OpenAI API
+key to capture.
+
+**How spans map to steps** (by ``span_data.type``):
+
+- ``generation`` / ``response`` → ``TURN``
+- ``function`` → ``TOOL`` (name, JSON ``input`` args, and the tool ``output``)
+- ``handoff`` → ``TOOL`` named ``"handoff"`` with ``{from, to}`` — control-flow, kept in-trajectory
+- ``guardrail`` → ``TOOL`` named ``"guardrail"`` with ``{name, triggered}`` — so a flip is diffable
+
+Spans are buffered per trace and, on ``on_trace_end``, sorted by ``started_at`` and mapped into one
+run — so parallel tool calls / concurrent agents (which interleave callbacks across threads) are
+handled. The initial input and final answer are recovered from the response spans, since the SDK
+exposes no trace-level input/output. Incomplete traces drain on ``force_flush`` / ``shutdown``.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from typing import Any, Optional
+
+from .. import recorder
+from ..errors import BrooderError
+from ..models import Step, StepKind
+from .base import parse_json
+
+# span_data.type values Brooder treats as model turns.
+_TURN_TYPES = frozenset({"generation", "response"})
+
+
+def _jsonable(value: Any) -> Any:
+    """Coerce a value into a JSON-safe form (unknown objects become their ``str``)."""
+    try:
+        return json.loads(json.dumps(value, default=str))
+    except (TypeError, ValueError):
+        return str(value)
+
+
+def _started_at(span: Any) -> str:
+    """Return a span's ISO start time for ordering (ISO strings sort chronologically)."""
+    return getattr(span, "started_at", None) or ""
+
+
+def _to_step(span: Any) -> Optional[Step]:
+    """Map one span to a trajectory step, or None for structural / unhandled spans."""
+    data = getattr(span, "span_data", None)
+    kind = getattr(data, "type", None)
+    if kind in _TURN_TYPES:
+        return Step(kind=StepKind.TURN, name="turn", args={"provider": "openai"})
+    if kind == "function":
+        return Step(
+            kind=StepKind.TOOL,
+            name=getattr(data, "name", None) or "tool",
+            args=parse_json(getattr(data, "input", None)),
+            result=_jsonable(getattr(data, "output", None)),
+        )
+    if kind == "handoff":
+        return Step(
+            kind=StepKind.TOOL,
+            name="handoff",
+            args={"from": getattr(data, "from_agent", None), "to": getattr(data, "to_agent", None)},
+        )
+    if kind == "guardrail":
+        return Step(
+            kind=StepKind.TOOL,
+            name="guardrail",
+            args={
+                "name": getattr(data, "name", None),
+                "triggered": bool(getattr(data, "triggered", False)),
+            },
+        )
+    return None
+
+
+def _initial_input(spans: list[Any]) -> Any:
+    """Recover the run's initial input from the earliest response/generation span."""
+    for span in spans:
+        data = getattr(span, "span_data", None)
+        if getattr(data, "type", None) in _TURN_TYPES:
+            value = getattr(data, "input", None)
+            if value is not None:
+                return _jsonable(value)
+    return None
+
+
+def _final_output(spans: list[Any]) -> Any:
+    """Recover the final answer from the latest response (``output_text``) or generation span."""
+    for span in reversed(spans):
+        data = getattr(span, "span_data", None)
+        kind = getattr(data, "type", None)
+        if kind == "response":
+            text = getattr(getattr(data, "response", None), "output_text", None)
+            if text:
+                return text
+        elif kind == "generation":
+            output = getattr(data, "output", None)
+            if output:
+                return _jsonable(output)
+    return None
+
+
+class BrooderTraceProcessor:
+    """An OpenAI Agents SDK ``TracingProcessor`` that ingests traces into Brooder trajectories.
+
+    Duck-typed against the SDK's ``TracingProcessor`` protocol, so importing this module does not
+    require ``openai-agents`` to be installed. Register it with :func:`install` (or add it yourself
+    via ``agents.add_trace_processor``).
+
+    Args:
+        agent: Logical agent name used to group baselines. Defaults to the trace's workflow name,
+            then ``"openai-agent"``.
+    """
+
+    def __init__(self, agent: Optional[str] = None) -> None:
+        """Create a processor recording into runs named ``agent`` (see the class docstring)."""
+        self._agent = agent
+        self._buffers: dict[str, dict[str, Any]] = {}
+        self._lock = threading.Lock()
+
+    def on_trace_start(self, trace: Any) -> None:
+        """Open a per-trace buffer keyed by ``trace_id``."""
+        trace_id = getattr(trace, "trace_id", None)
+        if trace_id is None:
+            return
+        with self._lock:
+            self._buffers.setdefault(trace_id, {"trace": trace, "spans": []})
+
+    def on_trace_end(self, trace: Any) -> None:
+        """Finalize the trace's buffered spans into one run."""
+        trace_id = getattr(trace, "trace_id", None)
+        if trace_id is not None:
+            self._finalize(trace_id)
+
+    def on_span_start(self, span: Any) -> None:
+        """Part of the protocol; capture happens at :meth:`on_span_end`."""
+        return None
+
+    def on_span_end(self, span: Any) -> None:
+        """Buffer an ended span under its trace id."""
+        trace_id = getattr(span, "trace_id", None)
+        if trace_id is None or trace_id == "no-op":
+            return
+        with self._lock:
+            entry = self._buffers.setdefault(trace_id, {"trace": None, "spans": []})
+            entry["spans"].append(span)
+
+    def shutdown(self, timeout: Optional[float] = None) -> None:
+        """Drain every buffered (possibly incomplete) trace into a run."""
+        for trace_id in list(self._buffers):
+            self._finalize(trace_id)
+
+    def force_flush(self) -> None:
+        """Drain every buffered trace into a run."""
+        for trace_id in list(self._buffers):
+            self._finalize(trace_id)
+
+    def _finalize(self, trace_id: str) -> None:
+        """Turn a trace's buffered spans into one recorded run."""
+        with self._lock:
+            entry = self._buffers.pop(trace_id, None)
+        if not entry or not entry["spans"]:
+            return
+        spans = entry["spans"]
+        spans.sort(key=_started_at)
+        agent = self._agent or getattr(entry["trace"], "name", None) or "openai-agent"
+        handle = recorder.open_run(agent, inputs=_initial_input(spans), external_id=trace_id)
+        for span in spans:
+            step = _to_step(span)
+            if step is not None:
+                handle.add_step(step)
+        handle.finish(_final_output(spans))
+
+
+def install(agent: Optional[str] = None) -> BrooderTraceProcessor:
+    """Register a Brooder trace processor with the OpenAI Agents SDK.
+
+    Appends the processor (the SDK's default exporter keeps running). Requires ``openai-agents``.
+
+    Args:
+        agent: Logical agent name used to group baselines (see :class:`BrooderTraceProcessor`).
+
+    Returns:
+        The registered :class:`BrooderTraceProcessor` (useful for tests / manual flushing).
+
+    Raises:
+        BrooderError: If the OpenAI Agents SDK is not installed.
+    """
+    try:
+        from agents import add_trace_processor
+    except ImportError as exc:
+        raise BrooderError(
+            "openai_agents.install() needs the OpenAI Agents SDK — `pip install openai-agents`"
+        ) from exc
+    processor = BrooderTraceProcessor(agent)
+    add_trace_processor(processor)
+    return processor