agentx-python 0.4.13__py3-none-any.whl → 0.5.1__py3-none-any.whl

agentx/integrations/__init__.py

@@ -5,3 +5,5 @@
 #   from agentx.integrations.crewai import AgentXCrewObserver
 #   from agentx.integrations.openai_agents import AgentXTracingProcessor
 #   from agentx.integrations.anthropic import patch_anthropic_client
+#   from agentx.integrations.google_adk import AgentXADKPlugin
+#   from agentx.integrations.google_genai import patch_genai_client

agentx/integrations/google_adk.py

@@ -0,0 +1,203 @@
+"""
+Google ADK integration for AgentX production tracing.
+
+Usage::
+
+    from agentx.integrations.google_adk import AgentXADKPlugin
+    from google.adk.runners import Runner
+    from google.adk.sessions import InMemorySessionService
+
+    plugin = AgentXADKPlugin(agentx.tracer, name="my-agent")
+
+    runner = Runner(
+        agent=agent,
+        app_name="my-app",
+        session_service=InMemorySessionService(),
+        plugins=[plugin],
+    )
+
+Requires: ``pip install "agentx-python[google-adk]"``
+"""
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, List, Optional
+
+from agentx.tracing.tracer import Tracer, _safe_serialize
+
+try:
+    from google.adk.plugins.base_plugin import BasePlugin
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        "google-adk is required for AgentXADKPlugin. "
+        "Install it with: pip install \"agentx-python[google-adk]\""
+    ) from exc
+
+
+def _content_to_text(content: Any) -> Optional[str]:
+    """Extract plain text from a google.genai types.Content object."""
+    if content is None:
+        return None
+    if isinstance(content, str):
+        return content
+    parts = getattr(content, "parts", None)
+    if not parts:
+        return None
+    texts = []
+    for part in parts:
+        text = getattr(part, "text", None)
+        if text and isinstance(text, str):
+            texts.append(text)
+    return " ".join(texts) if texts else None
+
+
+class AgentXADKPlugin(BasePlugin):
+    """
+    Google ADK plugin that sends one AgentX trace per runner invocation.
+
+    Captures input (user message), output (final model reply), model name,
+    tool calls, and latency via the ADK plugin callback hooks.
+
+    Register via the ``plugins`` list when constructing the ADK ``Runner``.
+    """
+
+    def __init__(
+        self,
+        tracer: Tracer,
+        name: str = "google-adk-agent",
+        metadata: Optional[Dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> None:
+        super().__init__(name="agentx")
+        self._tracer = tracer
+        self._agent_name = name
+        self._metadata = metadata
+        self._session_id = session_id
+        # invocation_id → accumulated run state
+        self._runs: Dict[str, Dict[str, Any]] = {}
+        # id(tool_context) → start time float
+        self._tool_starts: Dict[int, float] = {}
+
+    # ------------------------------------------------------------------
+    # Run lifecycle
+    # ------------------------------------------------------------------
+
+    async def before_run_callback(self, *, invocation_context: Any) -> None:
+        inv_id = invocation_context.invocation_id
+        agent_name = getattr(invocation_context.agent, "name", None) or self._agent_name
+        self._runs[inv_id] = {
+            "start": time.time(),
+            "name": agent_name,
+            "input": None,
+            "output": None,
+            "model": None,
+            "tool_calls": [],
+            "error": None,
+        }
+
+    async def on_user_message_callback(
+        self, *, invocation_context: Any, user_message: Any
+    ) -> None:
+        inv_id = invocation_context.invocation_id
+        state = self._runs.get(inv_id)
+        if state is None:
+            return
+        text = _content_to_text(user_message)
+        if text and state["input"] is None:
+            state["input"] = text
+
+    async def after_run_callback(self, *, invocation_context: Any) -> None:
+        inv_id = invocation_context.invocation_id
+        state = self._runs.pop(inv_id, None)
+        if state is None:
+            return
+        latency_ms = int((time.time() - state["start"]) * 1000)
+        self._tracer._send(
+            name=state["name"],
+            input=state["input"],
+            output=state["output"],
+            latency_ms=latency_ms,
+            framework="google-adk",
+            model=state["model"],
+            tool_calls=state["tool_calls"] or None,
+            metadata=self._metadata,
+            session_id=self._session_id,
+        )
+
+    # ------------------------------------------------------------------
+    # Model callbacks — capture model name and output
+    # ------------------------------------------------------------------
+
+    async def before_model_callback(
+        self, *, callback_context: Any, llm_request: Any
+    ) -> None:
+        inv_id = callback_context.get_invocation_context().invocation_id
+        state = self._runs.get(inv_id)
+        if state and not state["model"] and llm_request.model:
+            state["model"] = str(llm_request.model)
+
+    async def after_model_callback(
+        self, *, callback_context: Any, llm_response: Any
+    ) -> None:
+        inv_id = callback_context.get_invocation_context().invocation_id
+        state = self._runs.get(inv_id)
+        if state is None:
+            return
+        content = getattr(llm_response, "content", None)
+        text = _content_to_text(content)
+        # Keep updating — the last non-empty model reply is the final answer
+        if text:
+            state["output"] = text
+
+    # ------------------------------------------------------------------
+    # Tool callbacks
+    # ------------------------------------------------------------------
+
+    async def before_tool_callback(
+        self, *, tool: Any, tool_args: Dict[str, Any], tool_context: Any
+    ) -> None:
+        self._tool_starts[id(tool_context)] = time.time()
+
+    async def after_tool_callback(
+        self,
+        *,
+        tool: Any,
+        tool_args: Dict[str, Any],
+        tool_context: Any,
+        result: Dict[str, Any],
+    ) -> None:
+        inv_id = tool_context.get_invocation_context().invocation_id
+        state = self._runs.get(inv_id)
+        if state is None:
+            return
+        start = self._tool_starts.pop(id(tool_context), None)
+        tool_call: Dict[str, Any] = {
+            "name": getattr(tool, "name", "unknown"),
+            "input": _safe_serialize(tool_args),
+            "output": str(result)[:500] if result is not None else None,
+        }
+        if start is not None:
+            tool_call["latency_ms"] = max(0, int((time.time() - start) * 1000))
+        state["tool_calls"].append(tool_call)
+
+    async def on_tool_error_callback(
+        self,
+        *,
+        tool: Any,
+        tool_args: Dict[str, Any],
+        tool_context: Any,
+        error: Exception,
+    ) -> None:
+        inv_id = tool_context.get_invocation_context().invocation_id
+        state = self._runs.get(inv_id)
+        if state is None:
+            return
+        start = self._tool_starts.pop(id(tool_context), None)
+        tool_call: Dict[str, Any] = {
+            "name": getattr(tool, "name", "unknown"),
+            "input": _safe_serialize(tool_args),
+            "output": f"ERROR: {error}",
+        }
+        if start is not None:
+            tool_call["latency_ms"] = max(0, int((time.time() - start) * 1000))
+        state["tool_calls"].append(tool_call)

agentx/integrations/google_genai.py

@@ -0,0 +1,152 @@
+"""
+Google Gen AI SDK integration for AgentX production tracing.
+
+Usage::
+
+    from agentx.integrations.google_genai import patch_genai_client
+    from google import genai
+
+    client = genai.Client(api_key="GEMINI_API_KEY")
+    patch_genai_client(client, agentx.tracer, name="gemini-agent")
+
+    # All subsequent client.models.generate_content() calls are now traced.
+
+Requires: ``pip install "agentx-python[google-genai]"``
+"""
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, Optional
+
+from agentx.tracing.tracer import Tracer, _safe_serialize
+
+
+def patch_genai_client(
+    client: Any,
+    tracer: Tracer,
+    name: str = "gemini-agent",
+    metadata: Optional[Dict[str, Any]] = None,
+    session_id: Optional[str] = None,
+) -> None:
+    """
+    Monkey-patch ``client.models.generate_content`` (and
+    ``client.models.generate_content_stream`` if present) to automatically
+    send a trace for every call.
+
+    The original methods are still called and their return values are passed
+    through unchanged — nothing in the caller needs to change.
+
+    Calling this function on an already-patched client is a no-op.
+    """
+    models = getattr(client, "models", None)
+    if models is None:
+        raise ValueError("Provided client does not have a .models attribute")
+
+    _patch_generate_content(models, tracer, name, metadata, session_id)
+    if hasattr(models, "generate_content_stream"):
+        _patch_generate_content_stream(models, tracer, name, metadata, session_id)
+
+
+def _extract_response_text(response: Any) -> Optional[str]:
+    """Pull the generated text out of a GenerateContentResponse."""
+    # Convenience .text property (available on non-streaming responses)
+    text = getattr(response, "text", None)
+    if text and isinstance(text, str):
+        return text
+    # Fallback: walk candidates → content → parts
+    candidates = getattr(response, "candidates", None) or []
+    for candidate in candidates:
+        content = getattr(candidate, "content", None)
+        parts = getattr(content, "parts", None) or []
+        for part in parts:
+            t = getattr(part, "text", None)
+            if t and isinstance(t, str):
+                return t
+    return None
+
+
+def _patch_generate_content(
+    models: Any,
+    tracer: Tracer,
+    name: str,
+    metadata: Optional[Dict[str, Any]],
+    session_id: Optional[str],
+) -> None:
+    original = models.generate_content
+    if getattr(original, "_agentx_patched", False):
+        return
+
+    def patched(*args, **kwargs):
+        start = time.time()
+        error: Optional[str] = None
+        response = None
+        try:
+            response = original(*args, **kwargs)
+            return response
+        except Exception as exc:
+            error = str(exc)
+            raise
+        finally:
+            latency_ms = int((time.time() - start) * 1000)
+            model = kwargs.get("model") or (args[0] if args else None)
+            contents = kwargs.get("contents") or (args[1] if len(args) > 1 else None)
+            output = _extract_response_text(response) if response is not None else None
+            tracer._send(
+                name=name,
+                input=contents if isinstance(contents, str) else _safe_serialize(contents),
+                output=output,
+                latency_ms=latency_ms,
+                error=error,
+                framework="google-genai",
+                model=str(model) if model else None,
+                metadata=metadata,
+                session_id=session_id,
+            )
+
+    patched._agentx_patched = True
+    models.generate_content = patched
+
+
+def _patch_generate_content_stream(
+    models: Any,
+    tracer: Tracer,
+    name: str,
+    metadata: Optional[Dict[str, Any]],
+    session_id: Optional[str],
+) -> None:
+    original_stream = models.generate_content_stream
+    if getattr(original_stream, "_agentx_patched", False):
+        return
+
+    def patched_stream(*args, **kwargs):
+        start = time.time()
+        accumulated_text: list[str] = []
+        error: Optional[str] = None
+
+        try:
+            for chunk in original_stream(*args, **kwargs):
+                text = getattr(chunk, "text", None)
+                if text:
+                    accumulated_text.append(text)
+                yield chunk
+        except Exception as exc:
+            error = str(exc)
+            raise
+        finally:
+            latency_ms = int((time.time() - start) * 1000)
+            model = kwargs.get("model") or (args[0] if args else None)
+            contents = kwargs.get("contents") or (args[1] if len(args) > 1 else None)
+            tracer._send(
+                name=name,
+                input=contents if isinstance(contents, str) else _safe_serialize(contents),
+                output="".join(accumulated_text) or None,
+                latency_ms=latency_ms,
+                error=error,
+                framework="google-genai",
+                model=str(model) if model else None,
+                metadata=metadata,
+                session_id=session_id,
+            )
+
+    patched_stream._agentx_patched = True
+    models.generate_content_stream = patched_stream

agentx/integrations/langchain.py

@@ -10,10 +10,13 @@ Usage::
     # LCEL chain
     chain.invoke({"query": q}, config={"callbacks": [handler]})
 
+    # LangGraph agent (create_agent / create_react_agent)
+    agent.invoke({"messages": [...]}, config={"callbacks": [handler]})
+
     # AgentExecutor
     agent.invoke({"input": q}, config={"callbacks": [handler]})
 
-Requires: ``pip install agentx[langchain]``
+Requires: ``pip install "agentx-python[langchain]"``
 """
 from __future__ import annotations
 
@@ -29,14 +32,132 @@ try:
 except ImportError as exc:  # pragma: no cover
     raise ImportError(
         "langchain-core is required for AgentXCallbackHandler. "
-        "Install it with: pip install agentx[langchain]"
+        "Install it with: pip install \"agentx-python[langchain]\""
     ) from exc
 
 
+# ---------------------------------------------------------------------------
+# Output extraction helpers
+# ---------------------------------------------------------------------------
+
+def _extract_output(outputs: Any) -> Any:
+    """
+    Extract a clean, human-readable output from a chain's return value.
+
+    - LangGraph agents return {"messages": [HumanMessage, ..., AIMessage(final)]}
+      → extract the last AI message's text content
+    - AgentExecutor returns {"output": "..."}
+    - LCEL chains return {"text": "..."} or a plain string
+    """
+    if not isinstance(outputs, dict):
+        return _safe_serialize(outputs)
+
+    # LangGraph: {"messages": [...]}
+    messages = outputs.get("messages")
+    if isinstance(messages, list) and messages:
+        for msg in reversed(messages):
+            msg_type = getattr(msg, "type", None) or getattr(msg, "role", None)
+            if msg_type in ("ai", "assistant"):
+                content = getattr(msg, "content", None)
+                if content and isinstance(content, str) and content.strip():
+                    return content
+                if isinstance(content, list):
+                    # Multi-part content blocks
+                    texts = [p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"]
+                    joined = " ".join(t for t in texts if t).strip()
+                    if joined:
+                        return joined
+
+    # Standard chain / AgentExecutor
+    for key in ("output", "text", "answer", "result", "response"):
+        val = outputs.get(key)
+        if val and isinstance(val, str):
+            return val
+
+    return _safe_serialize(outputs)
+
+
+def _extract_input(inputs: Any) -> Any:
+    """
+    Extract a clean input from a chain's input dict.
+
+    - LangGraph: {"messages": [HumanMessage(...)]} → first human message text
+    - AgentExecutor: {"input": "..."} → the string
+    - LCEL: {"query": "...", "question": "..."} → the string value
+    """
+    if not isinstance(inputs, dict):
+        return _safe_serialize(inputs)
+
+    # LangGraph: {"messages": [...]}
+    messages = inputs.get("messages")
+    if isinstance(messages, list) and messages:
+        for msg in messages:
+            msg_type = getattr(msg, "type", None) or getattr(msg, "role", None)
+            if msg_type in ("human", "user"):
+                content = getattr(msg, "content", None)
+                if isinstance(content, str):
+                    return content
+        # Fallback: first message regardless of type
+        first = messages[0]
+        content = getattr(first, "content", None)
+        if isinstance(content, str):
+            return content
+
+    # Standard
+    for key in ("input", "query", "question", "human_input"):
+        val = inputs.get(key)
+        if val and isinstance(val, str):
+            return val
+
+    return _safe_serialize(inputs)
+
+
+def _extract_tool_calls_from_messages(outputs: Any) -> List[Dict[str, Any]]:
+    """
+    Fallback: extract tool calls directly from the message history when
+    on_tool_end callbacks were not linked to the top-level run.
+    """
+    if not isinstance(outputs, dict):
+        return []
+
+    messages = outputs.get("messages")
+    if not isinstance(messages, list):
+        return []
+
+    # Build map of tool_call_id → result from ToolMessages
+    results: Dict[str, str] = {}
+    for msg in messages:
+        if getattr(msg, "type", None) == "tool":
+            call_id = getattr(msg, "tool_call_id", None)
+            content = getattr(msg, "content", "")
+            if call_id:
+                results[call_id] = str(content)
+
+    # Extract tool calls from AIMessages
+    tool_calls: List[Dict[str, Any]] = []
+    for msg in messages:
+        if getattr(msg, "type", None) != "ai":
+            continue
+        for tc in getattr(msg, "tool_calls", []) or []:
+            if not isinstance(tc, dict):
+                continue
+            call_id = tc.get("id", "")
+            tool_calls.append({
+                "name": tc.get("name", "unknown"),
+                "input": str(tc.get("args", "")),
+                "output": results.get(call_id, ""),
+            })
+
+    return tool_calls
+
+
 class AgentXCallbackHandler(BaseCallbackHandler):
     """
     LangChain callback handler that captures the top-level chain run and all
     nested tool calls, then sends one trace per top-level chain invocation.
+
+    Compatible with AgentExecutor, LCEL chains, and LangGraph agents
+    (``create_agent``, ``create_react_agent``).
     """
 
     def __init__(
@@ -52,10 +173,10 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         self._metadata = metadata
         self._session_id = session_id
 
-        # Keyed by run_id (UUID) → state dict
         self._runs: Dict[UUID, Dict[str, Any]] = {}
-        # Track which run_ids are top-level (no parent)
         self._top_level: Dict[UUID, bool] = {}
+        # Full parent-chain map so _find_top_ancestor can walk arbitrary depth
+        self._parents: Dict[UUID, Optional[UUID]] = {}
 
     # ------------------------------------------------------------------
     # Chain lifecycle
@@ -70,12 +191,13 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs,
     ) -> None:
+        self._parents[run_id] = parent_run_id
         is_top = parent_run_id is None
         self._top_level[run_id] = is_top
         if is_top:
             self._runs[run_id] = {
                 "start": time.time(),
-                "input": _safe_serialize(inputs),
+                "input": _extract_input(inputs),
                 "tool_calls": [],
                 "model": None,
             }
@@ -94,18 +216,22 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         if state is None:
             return
         latency_ms = int((time.time() - state["start"]) * 1000)
+        output = _extract_output(outputs)
+        # If callbacks missed tool calls (deep nesting), extract from message history
+        tool_calls = state["tool_calls"] or _extract_tool_calls_from_messages(outputs)
         self._tracer._send(
             name=self._name,
             input=state["input"],
-            output=_safe_serialize(outputs),
+            output=output,
             latency_ms=latency_ms,
             framework="langchain",
             model=state.get("model"),
-            tool_calls=state["tool_calls"] or None,
+            tool_calls=tool_calls or None,
             metadata=self._metadata,
             session_id=self._session_id,
         )
         self._top_level.pop(run_id, None)
+        self._parents.pop(run_id, None)
 
     def on_chain_error(
         self,
@@ -133,9 +259,10 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             session_id=self._session_id,
         )
         self._top_level.pop(run_id, None)
+        self._parents.pop(run_id, None)
 
     # ------------------------------------------------------------------
-    # LLM lifecycle (captures model name and token counts)
+    # LLM lifecycle
     # ------------------------------------------------------------------
 
     def on_llm_start(
@@ -147,18 +274,21 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs,
     ) -> None:
-        # Record LLM start time keyed by its own run_id
+        self._parents[run_id] = parent_run_id
         self._runs.setdefault(run_id, {})["llm_start"] = time.time()
 
-        # Propagate model name to the top-level run
         top = self._find_top_ancestor(parent_run_id)
         if top and not self._runs[top].get("model"):
+            # Check all common field locations across langchain versions
+            kw = serialized.get("kwargs", {})
             model = (
-                serialized.get("kwargs", {}).get("model_name")
-                or serialized.get("kwargs", {}).get("model")
-                or serialized.get("id", [None])[-1]
+                kw.get("model_name")
+                or kw.get("model")
+                or kwargs.get("invocation_params", {}).get("model")
+                or kwargs.get("invocation_params", {}).get("model_name")
+                or serialized.get("name")
             )
-            if model:
+            if model and model not in ("None", "none"):
                 self._runs[top]["model"] = str(model)
 
     def on_llm_end(
@@ -170,6 +300,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         **kwargs,
     ) -> None:
         self._runs.pop(run_id, None)
+        self._parents.pop(run_id, None)
 
     # ------------------------------------------------------------------
     # Tool lifecycle
@@ -184,6 +315,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs,
     ) -> None:
+        self._parents[run_id] = parent_run_id
         self._runs[run_id] = {
             "tool_name": serialized.get("name", "unknown"),
             "tool_input": input_str,
@@ -199,6 +331,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         **kwargs,
     ) -> None:
         state = self._runs.pop(run_id, None)
+        self._parents.pop(run_id, None)
         if state is None:
             return
         latency_ms = int((time.time() - state["start"]) * 1000)
@@ -221,6 +354,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         **kwargs,
     ) -> None:
         state = self._runs.pop(run_id, None)
+        self._parents.pop(run_id, None)
         if state is None:
             return
         tool_call = {
@@ -237,13 +371,13 @@ class AgentXCallbackHandler(BaseCallbackHandler):
     # Helpers
     # ------------------------------------------------------------------
 
-    def _find_top_ancestor(self, parent_run_id: Optional[UUID]) -> Optional[UUID]:
-        """Walk up parent chain to find the top-level run_id."""
-        current = parent_run_id
-        while current is not None:
+    def _find_top_ancestor(self, start: Optional[UUID]) -> Optional[UUID]:
+        """Walk the full parent chain to find the top-level run_id."""
+        current = start
+        seen: set = set()
+        while current is not None and current not in seen:
+            seen.add(current)
             if self._top_level.get(current):
                 return current
-            # If current is a nested run, keep climbing — for simplicity return
-            # the immediate parent that is registered as top-level
-            break
-        return current if current in self._runs else None
+            current = self._parents.get(current)
+        return None

agentx/integrations/openai_agents.py

@@ -11,22 +11,111 @@ Usage::
     from agents import add_trace_processor
     add_trace_processor(processor)
 
-Requires: ``pip install agentx[openai-agents]``
+Requires: ``pip install "agentx-python[openai-agents]"``
 """
 from __future__ import annotations
 
 import time
-from typing import Any, Dict, Optional
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
 
 
+def _iso_to_ts(iso: Optional[str]) -> Optional[float]:
+    """Parse an ISO-8601 timestamp string to a Unix timestamp float."""
+    if not iso:
+        return None
+    try:
+        return datetime.fromisoformat(iso.replace("Z", "+00:00")).timestamp()
+    except Exception:
+        return None
+
+
+def _span_latency_ms(span: Any) -> Optional[int]:
+    """Return span duration in ms from ISO started_at / ended_at strings."""
+    t0 = _iso_to_ts(getattr(span, "started_at", None))
+    t1 = _iso_to_ts(getattr(span, "ended_at", None))
+    if t0 is not None and t1 is not None:
+        return max(0, int((t1 - t0) * 1000))
+    return None
+
+
+def _extract_text(output: Any) -> Optional[str]:
+    """
+    Best-effort extraction of a human-readable string from a GenerationSpanData
+    or ResponseSpanData output, which can be a list of message dicts in either
+    the Chat Completions or Responses API format.
+    """
+    if output is None:
+        return None
+    if isinstance(output, str):
+        return output
+    if not isinstance(output, (list, tuple)):
+        return _safe_serialize(output)
+
+    for item in reversed(output):
+        if not isinstance(item, dict):
+            # Could be a pydantic model — try .text or .content
+            text = getattr(item, "text", None) or getattr(item, "content", None)
+            if text and isinstance(text, str):
+                return text
+            continue
+
+        # Responses API: {"type": "message", "content": [{"type": "output_text", "text": "..."}]}
+        if item.get("type") == "message":
+            content = item.get("content", [])
+            if isinstance(content, list):
+                for part in content:
+                    if isinstance(part, dict) and part.get("type") == "output_text":
+                        return part.get("text")
+
+        # Chat Completions API: {"role": "assistant", "content": "..."}
+        role = item.get("role", "")
+        if role == "assistant":
+            content = item.get("content")
+            if isinstance(content, str):
+                return content
+            if isinstance(content, list):
+                for part in content:
+                    if isinstance(part, dict) and part.get("type") == "text":
+                        return part.get("text")
+
+    return _safe_serialize(output)
+
+
+def _extract_input_text(input_data: Any) -> Optional[str]:
+    """Extract the user's input query from a generation span's input messages."""
+    if input_data is None:
+        return None
+    if isinstance(input_data, str):
+        return input_data
+    if not isinstance(input_data, (list, tuple)):
+        return _safe_serialize(input_data)
+
+    # Walk messages and grab the last user message content
+    user_text = None
+    for item in input_data:
+        if not isinstance(item, dict):
+            continue
+        if item.get("role") == "user":
+            content = item.get("content")
+            if isinstance(content, str):
+                user_text = content
+            elif isinstance(content, list):
+                for part in content:
+                    if isinstance(part, dict) and part.get("type") in ("text", "input_text"):
+                        user_text = part.get("text")
+
+    return user_text or _safe_serialize(input_data)
+
+
 class AgentXTracingProcessor:
     """
     Implements the ``TracingProcessor`` interface expected by the OpenAI Agents SDK
     (``agents.add_trace_processor``).
 
-    Sends one AgentX trace per top-level agent run span.
+    Sends one AgentX trace per top-level agent run.
     """
 
     def __init__(
@@ -38,7 +127,7 @@ class AgentXTracingProcessor:
         self._tracer = tracer
         self._metadata = metadata
         self._session_id = session_id
-        # span_id → state for in-flight agent spans
+        # trace_id → accumulated state
         self._spans: Dict[str, Dict[str, Any]] = {}
 
     # ------------------------------------------------------------------
@@ -46,18 +135,17 @@ class AgentXTracingProcessor:
     # ------------------------------------------------------------------
 
     def on_trace_start(self, trace: Any) -> None:
-        """Called when a new trace (top-level run) begins."""
         trace_id = getattr(trace, "trace_id", None) or str(id(trace))
         self._spans[trace_id] = {
             "start": time.time(),
             "name": getattr(trace, "name", "openai-agent"),
             "input": None,
+            "output": None,
             "tool_calls": [],
             "model": None,
         }
 
     def on_trace_end(self, trace: Any) -> None:
-        """Called when the trace ends (agent run complete)."""
         trace_id = getattr(trace, "trace_id", None) or str(id(trace))
         state = self._spans.pop(trace_id, None)
         if state is None:
@@ -66,7 +154,7 @@ class AgentXTracingProcessor:
         self._tracer._send(
             name=state["name"],
             input=state.get("input"),
-            output=_safe_serialize(getattr(trace, "output", None)),
+            output=state.get("output"),
             latency_ms=latency_ms,
             framework="openai-agents",
             model=state.get("model"),
@@ -76,47 +164,58 @@ class AgentXTracingProcessor:
         )
 
     def on_span_start(self, span: Any) -> None:
-        """Called for each nested span (LLM call, tool call, etc.)."""
-        span_type = getattr(span, "span_data", None)
-        if span_type is None:
-            return
-
-        # Capture model from LLM generation spans
-        if hasattr(span_type, "model"):
-            trace_id = getattr(span, "trace_id", None)
-            if trace_id and trace_id in self._spans and not self._spans[trace_id].get("model"):
-                self._spans[trace_id]["model"] = str(span_type.model)
-
-        # Capture input on the root span
-        if hasattr(span_type, "input") and span_type.input:
-            trace_id = getattr(span, "trace_id", None)
-            if trace_id and trace_id in self._spans and not self._spans[trace_id].get("input"):
-                self._spans[trace_id]["input"] = _safe_serialize(span_type.input)
+        pass  # All data is captured in on_span_end when fields are fully populated
 
     def on_span_end(self, span: Any) -> None:
-        """Called when a nested span ends; captures tool-call results."""
         span_data = getattr(span, "span_data", None)
         if span_data is None:
             return
 
-        # Tool / function calls
-        if hasattr(span_data, "tool_name") or hasattr(span_data, "name"):
-            trace_id = getattr(span, "trace_id", None)
-            if trace_id and trace_id in self._spans:
-                tool_entry: Dict[str, Any] = {
-                    "name": getattr(span_data, "tool_name", None)
-                    or getattr(span_data, "name", "tool"),
-                    "input": _safe_serialize(getattr(span_data, "input", None)),
-                    "output": str(getattr(span_data, "output", ""))[:500],
-                }
-                if hasattr(span, "started_at") and hasattr(span, "ended_at"):
-                    try:
-                        tool_entry["latency_ms"] = int(
-                            (span.ended_at - span.started_at) * 1000
-                        )
-                    except Exception:
-                        pass
-                self._spans[trace_id]["tool_calls"].append(tool_entry)
+        trace_id = getattr(span, "trace_id", None)
+        if not trace_id or trace_id not in self._spans:
+            return
+
+        state = self._spans[trace_id]
+        span_type = getattr(span_data, "type", None)
+
+        if span_type == "generation":
+            # Capture input from the first generation span
+            if state["input"] is None and span_data.input:
+                state["input"] = _extract_input_text(span_data.input)
+            # Always update output to the latest generation (last one wins = final reply)
+            if span_data.output:
+                state["output"] = _extract_text(span_data.output)
+            # Capture model name
+            if not state["model"] and span_data.model:
+                state["model"] = str(span_data.model)
+
+        elif span_type == "response":
+            # Responses API path — extract from the response object
+            response = getattr(span_data, "response", None)
+            if response is not None:
+                if state["input"] is None:
+                    raw_input = getattr(span_data, "input", None)
+                    if raw_input:
+                        state["input"] = _extract_input_text(raw_input) if isinstance(raw_input, list) else str(raw_input)
+                output_items = getattr(response, "output", None)
+                if output_items:
+                    state["output"] = _extract_text(output_items)
+                if not state["model"]:
+                    model = getattr(response, "model", None)
+                    if model:
+                        state["model"] = str(model)
+
+        elif span_type == "function":
+            # Tool / function call
+            tool_entry: Dict[str, Any] = {
+                "name": span_data.name,
+                "input": span_data.input,
+                "output": str(span_data.output)[:500] if span_data.output is not None else None,
+            }
+            latency = _span_latency_ms(span)
+            if latency is not None:
+                tool_entry["latency_ms"] = latency
+            state["tool_calls"].append(tool_entry)
 
     def force_flush(self) -> None:
         self._tracer.flush()

agentx/version.py

@@ -1 +1 @@
-VERSION = "0.4.13"
+VERSION = "0.5.1"

{agentx_python-0.4.13.dist-info → agentx_python-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.4.13
+Version: 0.5.1
 Summary: Official Python SDK for AgentX (https://www.agentx.so/)
 Home-page: https://github.com/AgentX-ai/AgentX-python
 Author: Robin Wang and AgentX Team
@@ -23,11 +23,17 @@ Provides-Extra: openai-agents
 Requires-Dist: openai-agents>=0.0.3; extra == "openai-agents"
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.25.0; extra == "anthropic"
+Provides-Extra: google-adk
+Requires-Dist: google-adk>=1.0.0; extra == "google-adk"
+Provides-Extra: google-genai
+Requires-Dist: google-genai>=1.0.0; extra == "google-genai"
 Provides-Extra: all
 Requires-Dist: langchain-core>=0.1.0; extra == "all"
 Requires-Dist: crewai>=0.80.0; extra == "all"
 Requires-Dist: openai-agents>=0.0.3; extra == "all"
 Requires-Dist: anthropic>=0.25.0; extra == "all"
+Requires-Dist: google-adk>=1.0.0; extra == "all"
+Requires-Dist: google-genai>=1.0.0; extra == "all"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{agentx_python-0.4.13.dist-info → agentx_python-0.5.1.dist-info}/RECORD

@@ -2,7 +2,7 @@ agentx/__init__.py,sha256=4iMiGLNU4S-I-WyX0Z7l0BFUY_krVJ3seMQPqyO-ReA,602
 agentx/agentx.py,sha256=tw3CUnkvTM1eopoO7u8TgKnAz9t0HNSxrVcjIJsytAg,3469
 agentx/exceptions.py,sha256=TKooc1_pQ2P_4DqpLfZx_ap-6cfrjFLW8JS2qaGx_1M,1237
 agentx/util.py,sha256=yOV3VVdc0KUbw5pEkSRiO2lZzEtL1vOsnB_f5tFk948,665
-agentx/version.py,sha256=dO83YWKEtJW3A7sv1BcvYqTJgQZCIf7sMXiONZFOY_Q,19
+agentx/version.py,sha256=Rm_tMXSiQe1ocDOCUKKIk-I7ZEL21NXSJBaSKX1ooCs,18
 agentx/evaluations/__init__.py,sha256=OmdCHiBO3Qcsck0UkOyVqSvfRs6nXbSi-2FXzyRfsog,89
 agentx/evaluations/_term.py,sha256=S4blZgH66wHvkt2hR29h8JE5xnakBeD7dlZdUW9Wfoo,2065
 agentx/evaluations/client.py,sha256=3mC0uSBRk9uO4F3a7kIMM_UUsyypF99_97E5MSpYFzM,7736
@@ -17,11 +17,13 @@ agentx/evaluations/adapters/__init__.py,sha256=fK8Hx75usbiY03XUSmnLnSrwBXip2CQs2
 agentx/evaluations/adapters/http_endpoint.py,sha256=3YDwpf9u6u3Xj8kLhERmShBeGMviKHCA0Dhr0t6KMmA,2061
 agentx/evaluations/adapters/precomputed.py,sha256=BNNLQuzYFWyLSwsrMGv6MvRTroKY5Hz-Lm4F2RszESU,1210
 agentx/evaluations/adapters/raw.py,sha256=FkDq_mdf21-xt5EHnGyIGcS6VZxkEkFcmM9VTd-T5sA,1130
-agentx/integrations/__init__.py,sha256=ZKOwYsvHXo9ipJUv5rPF1DOA0H5Q5YBG0WRI1f0lv8s,414
+agentx/integrations/__init__.py,sha256=xpgFqfqdMqXSdCSucQoPRaewj9w5upSmkcaQtBKrtRE,545
 agentx/integrations/anthropic.py,sha256=VzRdeICYm_zfgNBy6PEu6bYSUa9iKuSoRX6HSivWQpM,4802
 agentx/integrations/crewai.py,sha256=o7tgFE1I2T7fxOc2lG7aY1nAAwihejV3y-EJII6z1j0,3488
-agentx/integrations/langchain.py,sha256=_W62CUnzqlzSgIcNim0xDGXc3WsPoWYLQSHpTD4Xp5E,7844
-agentx/integrations/openai_agents.py,sha256=XDN_tNfx1ztLUlwNHgFk6QSiYbUDLo-HPljrHjZoVWk,4668
+agentx/integrations/google_adk.py,sha256=270HRdOuHo_4AWULoN0oyJ9XlrBKTeb4Ipjk5EZxzrY,6938
+agentx/integrations/google_genai.py,sha256=fNKyoqegBYhxAeADYuRaI5K6-p1h6ntAdpBaNhFziQQ,5147
+agentx/integrations/langchain.py,sha256=hSy2aRi0QzVo0eDpIRRacZg7ulJG-XqY_tH7C5YvNd0,12848
+agentx/integrations/openai_agents.py,sha256=JscJELxs8sTj3toi9ixY5TOpH-lCVBtpgxYgahIaezE,8162
 agentx/resources/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agentx/resources/agent.py,sha256=ZKpxDYzJbKgOX4-W86U__b4gko7XlbwB4q6L8Az9uo0,2011
 agentx/resources/conversation.py,sha256=HK7gEGK6qIUz0KvG-ItUip08otWZdkP5S6qHiznvgZM,4443
@@ -30,8 +32,8 @@ agentx/tracing/__init__.py,sha256=l2mIRILE-wWrCD0fekgrIOQUEvtDObzdrkgpQUtXroU,40
 agentx/tracing/ci_types.py,sha256=zFVcZGvc1qbVDdOlEPQ1i5R6terqJHxCzFF4ZjouxxI,1423
 agentx/tracing/ingest_client.py,sha256=8aLabvWecgZsn0pkajZkI1it2TOk5boP0kC3o6TO52o,11487
 agentx/tracing/tracer.py,sha256=nJILoeQg6HV4SR05xJPy6EiSiZCpkp4yFBTrN1BO_1w,16951
-agentx_python-0.4.13.dist-info/licenses/LICENSE,sha256=6ZbiPNFmv3xBb44LGhAa3PZYK0ROAztsd5LRFZDlGFE,1074
-agentx_python-0.4.13.dist-info/METADATA,sha256=nDiVYt09aeyL89w3-ro9WfjsD7abgG0DdyujwsOdC5Q,6630
-agentx_python-0.4.13.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-agentx_python-0.4.13.dist-info/top_level.txt,sha256=s-q-HB9Gb_QdrZNacSeQyF_c25gQooMy7DlxzgLOHPk,7
-agentx_python-0.4.13.dist-info/RECORD,,
+agentx_python-0.5.1.dist-info/licenses/LICENSE,sha256=6ZbiPNFmv3xBb44LGhAa3PZYK0ROAztsd5LRFZDlGFE,1074
+agentx_python-0.5.1.dist-info/METADATA,sha256=gq6txrb4Xjdar5TAzk-aHMiON2zM-i7CA4njI8UFLK8,6901
+agentx_python-0.5.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+agentx_python-0.5.1.dist-info/top_level.txt,sha256=s-q-HB9Gb_QdrZNacSeQyF_c25gQooMy7DlxzgLOHPk,7
+agentx_python-0.5.1.dist-info/RECORD,,