agentos-python 0.1.0__py3-none-any.whl

agentos/integrations/llamaindex.py

@@ -0,0 +1,174 @@
+"""LlamaIndex integration — CallbackHandler for auto-capturing LLM and retriever events.
+
+Usage::
+
+    from agentos.integrations.llamaindex import AgentOSCallbackHandler
+    from llama_index.core import Settings
+
+    handler = AgentOSCallbackHandler(api_key="aos_...", agent_id="my-agent")
+    Settings.callback_manager.add_handler(handler)
+
+    # All LlamaIndex LLM calls and retrievals auto-captured
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger("agentos.integrations.llamaindex")
+
+
+class AgentOSCallbackHandler:
+    """LlamaIndex-compatible callback handler.
+
+    Implements the LlamaIndex BaseCallbackHandler interface via duck typing
+    (no import required). Captures LLM calls, retriever queries, and tool use.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        agent_id: str = "llamaindex-agent",
+        base_url: str = "https://api.agentos.dev",
+        capture_content: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        from agentos.client import AgentOS, get_client
+
+        if api_key:
+            self._client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0, **kwargs)
+        else:
+            self._client = get_client()
+
+        self._agent_id = agent_id
+        self._capture_content = capture_content
+        self._starts: dict[str, float] = {}
+        self._event_data: dict[str, dict[str, Any]] = {}
+
+    # Required by LlamaIndex BaseCallbackHandler
+    def on_event_start(
+        self,
+        event_type: str,
+        payload: dict[str, Any] | None = None,
+        event_id: str = "",
+        parent_id: str = "",
+        **kwargs: Any,
+    ) -> str:
+        """Called when a LlamaIndex event starts."""
+        self._starts[event_id] = time.monotonic()
+        self._event_data[event_id] = {
+            "event_type": event_type,
+            "payload": payload or {},
+            "parent_id": parent_id,
+        }
+        return event_id
+
+    def on_event_end(
+        self,
+        event_type: str,
+        payload: dict[str, Any] | None = None,
+        event_id: str = "",
+        **kwargs: Any,
+    ) -> None:
+        """Called when a LlamaIndex event ends."""
+        if self._client is None:
+            return
+
+        start = self._starts.pop(event_id, None)
+        start_data = self._event_data.pop(event_id, {})
+        duration_ms = (time.monotonic() - start) * 1000 if start else None
+        payload = payload or {}
+
+        try:
+            if event_type == "llm":
+                self._handle_llm_end(payload, start_data, duration_ms)
+            elif event_type == "retrieve":
+                self._handle_retrieve_end(payload, start_data, duration_ms)
+            elif event_type == "function_call":
+                self._handle_tool_end(payload, start_data, duration_ms)
+        except Exception:
+            logger.exception("Failed to capture LlamaIndex event")
+
+    def _handle_llm_end(
+        self,
+        payload: dict[str, Any],
+        start_data: dict[str, Any],
+        duration_ms: float | None,
+    ) -> None:
+        """Process an LLM completion event."""
+        response = payload.get("response", "")
+        model = payload.get("model", "unknown")
+
+        data: dict[str, Any] = {
+            "model": model,
+            "system": "llamaindex",
+        }
+
+        if duration_ms is not None:
+            data["duration_ms"] = duration_ms
+
+        # Token counts from payload
+        token_counts = payload.get("token_counts", {})
+        if token_counts:
+            data["input_tokens"] = token_counts.get("prompt_tokens")
+            data["output_tokens"] = token_counts.get("completion_tokens")
+
+        if self._capture_content and response:
+            data["output"] = [{"role": "assistant", "content": str(response)}]
+
+        self._client.llm_call(self._agent_id, **data)  # type: ignore[union-attr]
+
+    def _handle_retrieve_end(
+        self,
+        payload: dict[str, Any],
+        start_data: dict[str, Any],
+        duration_ms: float | None,
+    ) -> None:
+        """Process a retrieval event."""
+        nodes = payload.get("nodes", [])
+
+        self._client.retrieval_query(  # type: ignore[union-attr]
+            self._agent_id,
+            source="llamaindex-retriever",
+            results_count=len(nodes),
+            duration_ms=duration_ms,
+        )
+
+    def _handle_tool_end(
+        self,
+        payload: dict[str, Any],
+        start_data: dict[str, Any],
+        duration_ms: float | None,
+    ) -> None:
+        """Process a tool/function call event."""
+        tool_name = payload.get("function_call", {}).get("name", "unknown")
+
+        data: dict[str, Any] = {
+            "tool_name": tool_name,
+            "status": "success",
+        }
+        if duration_ms is not None:
+            data["duration_ms"] = duration_ms
+
+        self._client.tool_call(self._agent_id, **data)  # type: ignore[union-attr]
+
+    # Required interface methods
+    def start_trace(self, trace_id: str | None = None) -> None:
+        pass
+
+    def end_trace(
+        self,
+        trace_id: str | None = None,
+        trace_map: dict[str, list[str]] | None = None,
+    ) -> None:
+        pass
+
+    def flush(self) -> None:
+        if self._client:
+            self._client.flush()
+
+    def shutdown(self) -> None:
+        if self._client:
+            self._client.shutdown()

agentos/integrations/openai.py

@@ -0,0 +1,149 @@
+"""OpenAI SDK integration — wraps client to auto-capture LLM calls.
+
+Usage::
+
+    from agentos.integrations.openai import wrap_openai
+    import openai
+
+    client = openai.OpenAI()
+    client = wrap_openai(client, api_key="aos_...", agent_id="my-agent")
+
+    # All chat.completions.create() calls auto-traced
+    response = client.chat.completions.create(model="gpt-4o", messages=[...])
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger("agentos.integrations.openai")
+
+
+def wrap_openai(
+    openai_client: Any,
+    *,
+    api_key: str | None = None,
+    agent_id: str = "default",
+    base_url: str = "https://api.agentos.dev",
+    capture_content: bool = True,
+    **kwargs: Any,
+) -> Any:
+    """Wrap an OpenAI client to auto-capture LLM calls.
+
+    Returns the same client with patched ``chat.completions.create`` method.
+    """
+    from agentos.client import AgentOS, get_client
+
+    aos_client = None
+    if api_key:
+        aos_client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0, **kwargs)
+    else:
+        aos_client = get_client()
+
+    if aos_client is None:
+        logger.warning("No AgentOS client configured — wrap_openai is a no-op")
+        return openai_client
+
+    original_create = openai_client.chat.completions.create
+
+    def _extract_event_data(
+        response: Any,
+        model_arg: str,
+        messages: Any,
+        start: float,
+        call_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Extract event data from an OpenAI response."""
+        duration_ms = (time.monotonic() - start) * 1000
+
+        data: dict[str, Any] = {
+            "model": getattr(response, "model", model_arg),
+            "system": "openai",
+            "duration_ms": duration_ms,
+        }
+
+        # Usage
+        usage = getattr(response, "usage", None)
+        if usage:
+            data["input_tokens"] = getattr(usage, "prompt_tokens", None)
+            data["output_tokens"] = getattr(usage, "completion_tokens", None)
+            data["total_tokens"] = getattr(usage, "total_tokens", None)
+
+        # Finish reason
+        choices = getattr(response, "choices", [])
+        if choices:
+            fr = getattr(choices[0], "finish_reason", None)
+            if fr:
+                data["finish_reason"] = fr
+
+        # Response ID
+        resp_id = getattr(response, "id", None)
+        if resp_id:
+            data["response_id"] = resp_id
+
+        # Content (if capture enabled)
+        if capture_content:
+            if messages:
+                data["input"] = [
+                    {"role": m.get("role", ""), "content": m.get("content")}
+                    for m in messages
+                    if isinstance(m, dict)
+                ]
+            if choices:
+                msg = getattr(choices[0], "message", None)
+                if msg:
+                    data["output"] = [
+                        {
+                            "role": getattr(msg, "role", "assistant"),
+                            "content": getattr(msg, "content", None),
+                        }
+                    ]
+
+        # Temperature, max_tokens from call args
+        if "temperature" in call_kwargs:
+            data["temperature"] = call_kwargs["temperature"]
+        if "max_tokens" in call_kwargs:
+            data["max_tokens"] = call_kwargs["max_tokens"]
+
+        return data
+
+    def patched_create(*args: Any, **call_kwargs: Any) -> Any:
+        model_arg = call_kwargs.get("model", "unknown")
+        messages = call_kwargs.get("messages")
+        stream = call_kwargs.get("stream", False)
+
+        start = time.monotonic()
+        try:
+            response = original_create(*args, **call_kwargs)
+
+            if stream:
+                # For streaming, return the stream as-is
+                # (full streaming capture would require wrapping the iterator)
+                return response
+
+            data = _extract_event_data(response, model_arg, messages, start, call_kwargs)
+            try:
+                aos_client.llm_call(agent_id, **data)
+            except Exception:
+                logger.exception("Failed to capture LLM call")
+
+            return response
+
+        except Exception as exc:
+            duration_ms = (time.monotonic() - start) * 1000
+            try:
+                aos_client.llm_call(
+                    agent_id,
+                    model=model_arg,
+                    system="openai",
+                    duration_ms=duration_ms,
+                    error={"type": type(exc).__name__, "message": str(exc)},
+                )
+            except Exception:
+                logger.exception("Failed to capture LLM error")
+            raise
+
+    openai_client.chat.completions.create = patched_create
+    return openai_client

agentos/integrations/pydantic_ai.py

@@ -0,0 +1,236 @@
+"""Pydantic AI integration — instruments Agent.run() to capture LLM calls and tool use.
+
+Usage::
+
+    from agentos.integrations.pydantic_ai import instrument_pydantic_ai
+
+    instrument_pydantic_ai(api_key="aos_...", agent_id="my-agent")
+
+    # All pydantic_ai Agent.run() calls auto-captured
+    from pydantic_ai import Agent
+    agent = Agent("openai:gpt-4o")
+    result = agent.run_sync("What is 2+2?")
+
+Or wrap a single agent::
+
+    from agentos.integrations.pydantic_ai import wrap_agent
+
+    agent = Agent("openai:gpt-4o")
+    agent = wrap_agent(agent, api_key="aos_...", agent_id="math-bot")
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger("agentos.integrations.pydantic_ai")
+
+
+def wrap_agent(
+    agent: Any,
+    *,
+    api_key: str | None = None,
+    agent_id: str = "pydantic-ai",
+    base_url: str = "https://api.agentos.dev",
+    capture_content: bool = True,
+) -> Any:
+    """Wrap a Pydantic AI Agent to auto-capture runs.
+
+    Patches ``run()``, ``run_sync()``, and ``run_stream()`` on the agent instance.
+    """
+    from agentos.client import AgentOS, get_client
+    from agentos.tracing import trace
+
+    client = get_client()
+    if client is None and api_key:
+        client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0)
+
+    if client is None:
+        logger.warning("No AgentOS client — wrap_agent is a no-op")
+        return agent
+
+    # Extract model info from agent
+    model_name = "unknown"
+    model_obj = getattr(agent, "model", None)
+    if model_obj is not None:
+        model_name = str(model_obj)
+
+    # Determine system from model string (e.g., "openai:gpt-4o" → "openai")
+    system = "unknown"
+    if ":" in model_name:
+        system = model_name.split(":")[0]
+
+    # Patch run_sync
+    if hasattr(agent, "run_sync"):
+        original_run_sync = agent.run_sync
+
+        @functools.wraps(original_run_sync)
+        def patched_run_sync(*args: Any, **kwargs: Any) -> Any:
+            start = time.monotonic()
+            with trace(agent_id=agent_id) as ctx:
+                try:
+                    result = original_run_sync(*args, **kwargs)
+                    duration_ms = (time.monotonic() - start) * 1000
+
+                    data_: dict[str, Any] = {
+                        "model": model_name,
+                        "system": system,
+                        "duration_ms": duration_ms,
+                        "trace_id": ctx.trace_id,
+                        "span_id": ctx.span_id,
+                    }
+
+                    # Extract usage from result if available
+                    usage = getattr(result, "usage", None) or getattr(result, "_usage", None)
+                    if usage:
+                        data_["input_tokens"] = getattr(usage, "request_tokens", None) or getattr(
+                            usage, "input_tokens", None
+                        )
+                        data_["output_tokens"] = getattr(usage, "response_tokens", None) or getattr(
+                            usage, "output_tokens", None
+                        )
+
+                    if capture_content:
+                        prompt = args[0] if args else kwargs.get("user_prompt")
+                        if prompt:
+                            data_["input"] = [{"role": "user", "content": str(prompt)}]
+                        result_data = getattr(result, "data", None)
+                        if result_data is not None:
+                            data_["output"] = [{"role": "assistant", "content": str(result_data)}]
+
+                    client.llm_call(agent_id, **data_)
+                    return result
+
+                except Exception as exc:
+                    duration_ms = (time.monotonic() - start) * 1000
+                    client.llm_call(
+                        agent_id,
+                        model=model_name,
+                        system=system,
+                        duration_ms=duration_ms,
+                        error={"type": type(exc).__name__, "message": str(exc)},
+                        trace_id=ctx.trace_id,
+                        span_id=ctx.span_id,
+                    )
+                    raise
+
+        agent.run_sync = patched_run_sync
+
+    # Patch async run
+    if hasattr(agent, "run"):
+        original_run = agent.run
+
+        @functools.wraps(original_run)
+        async def patched_run(*args: Any, **kwargs: Any) -> Any:
+            start = time.monotonic()
+            with trace(agent_id=agent_id) as ctx:
+                try:
+                    result = await original_run(*args, **kwargs)
+                    duration_ms = (time.monotonic() - start) * 1000
+
+                    data_: dict[str, Any] = {
+                        "model": model_name,
+                        "system": system,
+                        "duration_ms": duration_ms,
+                        "trace_id": ctx.trace_id,
+                        "span_id": ctx.span_id,
+                    }
+
+                    if capture_content:
+                        prompt = args[0] if args else kwargs.get("user_prompt")
+                        if prompt:
+                            data_["input"] = [{"role": "user", "content": str(prompt)}]
+                        result_data = getattr(result, "data", None)
+                        if result_data is not None:
+                            data_["output"] = [{"role": "assistant", "content": str(result_data)}]
+
+                    client.llm_call(agent_id, **data_)
+                    return result
+
+                except Exception as exc:
+                    duration_ms = (time.monotonic() - start) * 1000
+                    client.llm_call(
+                        agent_id,
+                        model=model_name,
+                        system=system,
+                        duration_ms=duration_ms,
+                        error={"type": type(exc).__name__, "message": str(exc)},
+                        trace_id=ctx.trace_id,
+                        span_id=ctx.span_id,
+                    )
+                    raise
+
+        agent.run = patched_run
+
+    return agent
+
+
+def instrument_pydantic_ai(
+    *,
+    api_key: str | None = None,
+    agent_id: str = "pydantic-ai",
+    base_url: str = "https://api.agentos.dev",
+    capture_content: bool = True,
+) -> None:
+    """Monkey-patch Pydantic AI's Agent class to auto-capture all runs.
+
+    Every ``Agent.run()`` and ``Agent.run_sync()`` call will be instrumented.
+    """
+    try:
+        from pydantic_ai import Agent
+    except ImportError:
+        logger.error("pydantic-ai not installed — cannot instrument")
+        return
+
+    from agentos.client import AgentOS, get_client
+
+    client = get_client()
+    if client is None and api_key:
+        client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0)
+
+    if client is None:
+        logger.warning("No AgentOS client — instrument_pydantic_ai is a no-op")
+        return
+
+    original_run_sync = Agent.run_sync
+
+    @functools.wraps(original_run_sync)
+    def patched_run_sync(self: Any, *args: Any, **kwargs: Any) -> Any:
+        from agentos.tracing import trace
+
+        model_name = str(getattr(self, "model", "unknown"))
+        system = model_name.split(":")[0] if ":" in model_name else "unknown"
+        name = getattr(self, "name", None) or agent_id
+
+        start = time.monotonic()
+        with trace(agent_id=name) as ctx:
+            try:
+                result = original_run_sync(self, *args, **kwargs)
+                duration_ms = (time.monotonic() - start) * 1000
+                client.llm_call(
+                    name,
+                    model=model_name,
+                    system=system,
+                    duration_ms=duration_ms,
+                    trace_id=ctx.trace_id,
+                    span_id=ctx.span_id,
+                )
+                return result
+            except Exception as exc:
+                duration_ms = (time.monotonic() - start) * 1000
+                client.llm_call(
+                    name,
+                    model=model_name,
+                    system=system,
+                    duration_ms=duration_ms,
+                    error={"type": type(exc).__name__, "message": str(exc)},
+                    trace_id=ctx.trace_id,
+                    span_id=ctx.span_id,
+                )
+                raise
+
+    Agent.run_sync = patched_run_sync  # type: ignore[assignment]
+    logger.info("Pydantic AI instrumented — Agent.run_sync patched")