infralo 0.1.0.dev1__py3-none-any.whl

infralo/integrations/agno.py

@@ -0,0 +1,381 @@
+"""
+infralo — Agno agent framework integration.
+
+Install
+-------
+    pip install infralo[agno]
+
+Quick start
+-----------
+    from infralo import Infralo
+    from infralo.integrations.agno import make_agno_callbacks
+    from agno.agent import Agent
+    from agno.models.openai import OpenAIChat
+
+    infralo = Infralo(api_key="vk_...")
+    callbacks = make_agno_callbacks(infralo)
+
+    with infralo.start_trace(session_id="user-123") as trace:
+        agent = Agent(
+            # Point at the Infralo gateway for full LLM+tool linkage
+            model=OpenAIChat(id="my-deployment", base_url="http://localhost:8000/v1"),
+            tools=[search_web, get_weather],
+            **callbacks,          # ← injects all four lifecycle hooks
+        )
+        agent.print_response("What is the weather in London?")
+
+What is captured
+----------------
+* Every tool invocation → a ``ToolSpan`` with input args, output, timing, status.
+* LLM→tool parent links when the model routes through the Infralo gateway
+  (``x-infralo-span-id`` response header is read from the raw HTTP response).
+* Model metadata (model name, token counts) is attached to the span as metadata
+  when exposed by Agno's ``ModelResponse``.
+
+LLM span linking
+----------------
+``after_model_callback`` walks several attribute paths to find the raw HTTP
+response headers returned by the gateway.  If the Agno version does not expose
+them, tool spans are still captured but will appear without a parent LLM span.
+Use ``trace.set_parent_span_id(span_id)`` to set it manually when needed.
+
+Thread / async safety
+---------------------
+Agno runs tools in async tasks that inherit the active ``ContextVar`` from the
+parent coroutine, so ``get_active_trace()`` always returns the correct trace.
+For sync Agno runs the ContextVar is still set on the calling thread.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+import uuid
+from typing import TYPE_CHECKING, Any
+
+try:
+    import agno
+except ImportError:
+    agno = None
+
+from infralo.trace import _patch_httpx_once, get_active_trace
+
+if TYPE_CHECKING:
+    from infralo.client import Infralo  # noqa: F401
+
+logger = logging.getLogger("infralo.integrations.agno")
+
+# Key used to stash the in-flight ToolSpan on tool_context.session_state
+# so the after_tool callback can retrieve and finalise it.
+_SPAN_KEY = "__infralo_tool_span__"
+
+
+def make_agno_callbacks(infralo: "Infralo") -> dict[str, Any]:  # noqa: ARG001
+    """
+    Return a dict of Agno lifecycle callbacks wired to the active Infralo trace.
+
+    Pass the dict as keyword arguments to ``Agent(...)``:
+
+        callbacks = make_agno_callbacks(infralo)
+
+        with infralo.start_trace(session_id="s1") as trace:
+            agent = Agent(model=..., tools=[...], **callbacks)
+            agent.print_response("...")
+
+    Returned keys:
+        - For Agno 2.x: ``tool_hooks`` and ``post_hooks``.
+        - For Agno 1.x: ``before_tool_callback``, ``after_tool_callback``,
+          and ``after_model_callback``.
+
+    Args:
+        infralo: Your ``Infralo`` client instance.  Not used at callback
+                 call time (trace is read from the ``ContextVar`` instead),
+                 but required for documentation parity with other integrations.
+    """
+    _patch_httpx_once()
+
+    is_agno_v2 = False
+    if agno is not None:
+        version = getattr(agno, "__version__", "1.0.0")
+        if version.split(".")[0] >= "2":
+            is_agno_v2 = True
+
+    if is_agno_v2:
+
+        class InfraloToolHooksList:
+            def __init__(self, sync_hook, async_hook):
+                self.sync_hook = sync_hook
+                self.async_hook = async_hook
+
+            def __iter__(self):
+                entrypoint_is_async = False
+                try:
+                    for frame_info in inspect.stack():
+                        frame_self = frame_info.frame.f_locals.get("self")
+                        if (
+                            frame_self is not None
+                            and frame_self.__class__.__name__ == "FunctionCall"
+                        ):
+                            entrypoint = getattr(
+                                getattr(frame_self, "function", None), "entrypoint", None
+                            )
+                            if entrypoint is not None:
+                                if inspect.iscoroutinefunction(entrypoint):
+                                    entrypoint_is_async = True
+                                break
+                except Exception:
+                    pass
+
+                if entrypoint_is_async:
+                    yield self.async_hook
+                else:
+                    yield self.sync_hook
+
+            def __reversed__(self):
+                return reversed(list(self))
+
+        def pre_hook(run_input: Any, run_context: Any) -> None:  # noqa: ARG001
+            trace = get_active_trace()
+            if (
+                trace is not None
+                and run_context is not None
+                and hasattr(run_context, "session_state")
+            ):
+                run_context.session_state["__infralo_trace__"] = trace
+            return None
+
+        async def async_tool_hook(
+            function_name: str, function_call: Any, arguments: dict, run_context: Any = None
+        ) -> Any:
+            trace = None
+            if run_context is not None and hasattr(run_context, "session_state"):
+                trace = run_context.session_state.get("__infralo_trace__")
+            if trace is None:
+                trace = get_active_trace()
+
+            if trace is None:
+                return await function_call(**arguments)
+
+            tool_call_id = str(uuid.uuid4())
+            with trace.tool(name=function_name, tool_call_id=tool_call_id) as span:
+                span.set_input(arguments)
+                span.set_metadata(framework="agno", version="2.x")
+                try:
+                    result = await function_call(**arguments)
+                    if hasattr(result, "content"):
+                        output = {"result": str(result.content)}
+                    elif isinstance(result, dict):
+                        output = result
+                    else:
+                        output = {"result": str(result)}
+                    span.set_output(output)
+                    return result
+                except Exception as e:
+                    raise e
+
+        def sync_tool_hook(
+            function_name: str, function_call: Any, arguments: dict, run_context: Any = None
+        ) -> Any:
+            trace = None
+            if run_context is not None and hasattr(run_context, "session_state"):
+                trace = run_context.session_state.get("__infralo_trace__")
+            if trace is None:
+                trace = get_active_trace()
+
+            if trace is None:
+                return function_call(**arguments)
+
+            tool_call_id = str(uuid.uuid4())
+            with trace.tool(name=function_name, tool_call_id=tool_call_id) as span:
+                span.set_input(arguments)
+                span.set_metadata(framework="agno", version="2.x")
+                try:
+                    result = function_call(**arguments)
+                    if hasattr(result, "content"):
+                        output = {"result": str(result.content)}
+                    elif isinstance(result, dict):
+                        output = result
+                    else:
+                        output = {"result": str(result)}
+                    span.set_output(output)
+                    return result
+                except Exception as e:
+                    raise e
+
+        def post_hook(run_output: Any, run_context: Any) -> None:  # noqa: ARG001
+            trace = None
+            if run_context is not None and hasattr(run_context, "session_state"):
+                trace = run_context.session_state.get("__infralo_trace__")
+            if trace is None:
+                trace = get_active_trace()
+
+            if trace is None:
+                return None
+
+            span_id = None
+            try:
+                for msg in getattr(run_output, "messages", []):
+                    if hasattr(msg, "response_metadata") and isinstance(
+                        msg.response_metadata, dict
+                    ):
+                        span_id = msg.response_metadata.get("x-infralo-span-id")
+                    if (
+                        not span_id
+                        and hasattr(msg, "provider_data")
+                        and isinstance(msg.provider_data, dict)
+                    ):
+                        span_id = msg.provider_data.get("x-infralo-span-id")
+            except Exception:
+                pass
+
+            if span_id:
+                trace.set_parent_span_id(span_id)
+            return None
+
+        return {
+            "pre_hooks": [pre_hook],
+            "tool_hooks": InfraloToolHooksList(sync_tool_hook, async_tool_hook),
+            "post_hooks": [post_hook],
+        }
+
+    # ── Tool callbacks (Agno 1.x / Phidata) ───────────────────────────────────
+
+    def before_tool_callback(tool: Any, args: dict, tool_context: Any) -> None:
+        """
+        Called by Agno before each tool function runs.
+        Starts a ToolSpan and stashes it on ``tool_context.session_state``.
+        Returns ``None`` so Agno proceeds with actual tool execution.
+        """
+        trace = get_active_trace()
+        if trace is None:
+            logger.warning(
+                "infralo: before_tool_callback fired outside an active trace. "
+                "Wrap your agent run with `with infralo.start_trace(): ...`."
+            )
+            return None
+
+        # Agno's ToolExecutionContext exposes tool_call_id directly.
+        tool_call_id: str = getattr(tool_context, "tool_call_id", None) or str(uuid.uuid4())
+        tool_name: str = str(
+            getattr(tool, "name", None) or getattr(tool, "__name__", None) or "unknown_tool"
+        )
+
+        try:
+            span = trace.start_tool_span(name=tool_name, tool_call_id=tool_call_id)
+            if args:
+                span.set_input(args)
+            # Attach framework metadata to the span
+            agent_name = getattr(tool_context, "agent_name", None)
+            session_id = getattr(tool_context, "session_id", None)
+            meta: dict[str, Any] = {"framework": "agno"}
+            if agent_name:
+                meta["agent_name"] = agent_name
+            if session_id:
+                meta["framework_session_id"] = session_id
+            span.set_metadata(**meta)
+
+            # Stash for after_tool_callback
+            session_state = getattr(tool_context, "session_state", None)
+            if session_state is not None:
+                session_state[_SPAN_KEY] = span
+            else:
+                # Fallback: attach directly (Agno may use different attr names)
+                tool_context.__dict__[_SPAN_KEY] = span
+        except Exception:
+            logger.debug("infralo: error starting tool span for '%s'", tool_name, exc_info=True)
+
+        return None  # always let Agno proceed with tool execution
+
+    def after_tool_callback(
+        tool: Any, args: dict, tool_context: Any, tool_response: Any  # noqa: ARG001
+    ) -> None:
+        """
+        Called by Agno after each tool function completes.
+        Finalises the ToolSpan with output and timing.
+        Returns ``None`` so Agno passes the original tool_response to the LLM.
+        """
+        # Retrieve the stashed span
+        span = None
+        session_state = getattr(tool_context, "session_state", None)
+        if session_state is not None:
+            span = session_state.pop(_SPAN_KEY, None)
+        if span is None:
+            span = tool_context.__dict__.pop(_SPAN_KEY, None)
+
+        if span is None:
+            return None
+
+        try:
+            # Agno's tool_response can be a ToolResponse, string, dict, etc.
+            if hasattr(tool_response, "content"):
+                output: Any = {"result": str(tool_response.content)}
+            elif isinstance(tool_response, dict):
+                output = tool_response
+            else:
+                output = {"result": str(tool_response)}
+            span.end(output=output)
+        except Exception:
+            logger.debug("infralo: error ending tool span", exc_info=True)
+            try:
+                span.end()
+            except Exception:
+                pass
+
+        return None  # pass response to LLM unchanged
+
+    # ── Model callbacks ─────────────────────────────────────────────────────
+
+    def after_model_callback(
+        agent: Any, messages: Any, model_response: Any  # noqa: ARG001
+    ) -> None:
+        """
+        Called by Agno after each LLM response.
+
+        Attempts to extract the ``x-infralo-span-id`` header from the raw HTTP
+        response (available when routing through the Infralo gateway) and sets it
+        as the current parent span so subsequent tool spans are correctly linked.
+
+        If the header is unavailable, a debug message is logged and tool spans
+        remain unlinked from the LLM span.
+        """
+        trace = get_active_trace()
+        if trace is None:
+            return None
+
+        span_id: str | None = None
+
+        # Walk possible attribute paths to find raw HTTP response headers.
+        # Agno stores the underlying OpenAI response in model_response.response.
+        # Different Agno/OpenAI SDK versions may expose headers via different paths.
+        try:
+            raw = getattr(model_response, "response", None)
+            if raw is not None:
+                headers = (
+                    getattr(raw, "headers", None)
+                    # openai >= 1.x wraps in APIResponse, exposed via _response
+                    or getattr(getattr(raw, "_response", None), "headers", None)
+                    # httpx-level response
+                    or getattr(getattr(raw, "http_response", None), "headers", None)
+                )
+                if headers:
+                    span_id = headers.get("x-infralo-span-id")
+        except Exception:
+            logger.debug("infralo: error reading model response headers", exc_info=True)
+
+        if span_id:
+            trace.set_parent_span_id(span_id)
+            logger.debug("infralo: set parent span id from Agno model response: %s", span_id)
+        else:
+            logger.debug(
+                "infralo: x-infralo-span-id not found in Agno model response. "
+                "Tool spans will be recorded without a parent LLM span. "
+                "Ensure the model base_url points to the Infralo gateway for full linkage."
+            )
+
+        return None  # return None to pass model_response through unchanged
+
+    return {
+        "before_tool_callback": before_tool_callback,
+        "after_tool_callback": after_tool_callback,
+        "after_model_callback": after_model_callback,
+    }