agentx-python 0.5.2__tar.gz → 0.6.1__tar.gz

{agentx_python-0.5.2/agentx_python.egg-info → agentx_python-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.5.2
+Version: 0.6.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

agentx_python-0.6.1/agentx/integrations/_perf.py

@@ -0,0 +1,116 @@
+"""
+Shared helper for building a performance_summary that matches the AgentX
+platform agent format.
+
+External framework integrations collect:
+  - execution_steps  — LLM inference calls (name, duration_ms, start_time, end_time)
+  - tool_call_steps  — function/tool invocations (same shape)
+
+``build_performance_summary`` merges them into the full structure expected by
+the backend ingest endpoint and the AgentX UI.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+
+def build_performance_summary(
+    total_duration_ms: float,
+    execution_steps: Optional[List[Dict[str, Any]]] = None,
+    tool_call_steps: Optional[List[Dict[str, Any]]] = None,
+    retrieval_steps: Optional[List[Dict[str, Any]]] = None,
+    has_errors: bool = False,
+) -> Dict[str, Any]:
+    """
+    Return a ``performance_summary`` dict compatible with the AgentX platform
+    agent format.
+
+    Args:
+        total_duration_ms:  Wall-clock duration of the entire run.
+        execution_steps:    LLM call entries — each must have ``name``,
+                            ``duration_ms``, and optionally ``start_time`` /
+                            ``end_time`` as Unix floats.
+        tool_call_steps:    Tool / function call entries — same shape.
+        retrieval_steps:    RAG / vector-store retrieval entries — same shape,
+                            plus optional ``query`` (str) and ``doc_count`` (int).
+        has_errors:         Set to True when the run ended with an exception.
+    """
+    execution_steps = execution_steps or []
+    tool_call_steps = tool_call_steps or []
+    retrieval_steps = retrieval_steps or []
+
+    # Merge all items with their phase type, sort by start_time
+    all_items: List[tuple[Dict[str, Any], str]] = (
+        [(s, "execution_step") for s in execution_steps]
+        + [(t, "tool_call") for t in tool_call_steps]
+        + [(r, "retrieval") for r in retrieval_steps]
+    )
+    all_items.sort(key=lambda x: x[0].get("start_time") or 0)
+
+    unified: List[Dict[str, Any]] = []
+    out_steps: List[Dict[str, Any]] = []
+    out_tools: List[Dict[str, Any]] = []
+    out_retrievals: List[Dict[str, Any]] = []
+
+    for order, (item, phase_type) in enumerate(all_items, start=1):
+        entry: Dict[str, Any] = {
+            "name": item["name"],
+            "duration_ms": round(float(item["duration_ms"]), 3),
+            "start_order": order,
+            "phase_type": phase_type,
+        }
+        if item.get("start_time") is not None:
+            entry["start_time"] = item["start_time"]
+        if item.get("end_time") is not None:
+            entry["end_time"] = item["end_time"]
+        # Retrieval-specific fields
+        if phase_type == "retrieval":
+            if item.get("query"):
+                entry["query"] = item["query"]
+            if item.get("doc_count") is not None:
+                entry["doc_count"] = item["doc_count"]
+        unified.append(entry)
+
+        flat = {k: v for k, v in entry.items() if k != "phase_type"}
+        if phase_type == "execution_step":
+            out_steps.append(flat)
+        elif phase_type == "tool_call":
+            out_tools.append(flat)
+        else:
+            out_retrievals.append(flat)
+
+    tools_total_ms = round(sum(t["duration_ms"] for t in out_tools), 3)
+    retrievals_total_ms = round(sum(r["duration_ms"] for r in out_retrievals), 3)
+
+    return {
+        "total_duration_ms": round(float(total_duration_ms), 3),
+        "todo_tasks_enabled": None,
+        "main_phases": [],
+        "tool_calls": out_tools,
+        "mcp_tool_calls": [],
+        "action_tool_calls": [],
+        "delegate_calls": [],
+        "execution_steps": out_steps,
+        "knowledge_retrievals": out_retrievals,
+        "todo_operations": [],
+        "detailed_phases": [],
+        "active_phases": [],
+        "has_errors": has_errors,
+        "memory_actions": None,
+        "unified_timeline": unified,
+        "statistics": {
+            "total_main_phases": 0,
+            "total_tool_calls": len(out_tools),
+            "total_mcp_tool_calls": 0,
+            "total_action_tool_calls": 0,
+            "total_delegate_calls": 0,
+            "total_execution_steps": len(out_steps),
+            "total_knowledge_retrievals": len(out_retrievals),
+            "total_todo_operations": 0,
+            "main_phases_total_ms": 0,
+            "tool_calls_total_ms": tools_total_ms,
+            "knowledge_retrievals_total_ms": retrievals_total_ms,
+            "mcp_tool_calls_total_ms": 0,
+            "todo_operations_total_ms": 0,
+        },
+    }

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx/integrations/anthropic.py

@@ -19,6 +19,7 @@ import time
 from typing import Any, Dict, Optional
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
+from agentx.integrations._perf import build_performance_summary
 
 
 def patch_anthropic_client(
@@ -58,7 +59,7 @@ def _patch_create(
         return  # already patched
 
     def patched_create(*args, **kwargs):
-        start = time.time()
+        start_t = time.time()
         error: Optional[str] = None
         response = None
         try:
@@ -68,15 +69,35 @@ def _patch_create(
             error = str(exc)
             raise
         finally:
-            latency_ms = int((time.time() - start) * 1000)
+            end_t = time.time()
+            latency_ms = int((end_t - start_t) * 1000)
             input_messages = kwargs.get("messages") or (args[0] if args else None)
             model = kwargs.get("model")
             output = None
+            input_tokens = None
+            output_tokens = None
             if response is not None:
                 try:
                     output = response.content[0].text if response.content else None
                 except Exception:
                     output = str(response)[:500]
+                try:
+                    usage = getattr(response, "usage", None)
+                    if usage is not None:
+                        input_tokens = getattr(usage, "input_tokens", None)
+                        output_tokens = getattr(usage, "output_tokens", None)
+                except Exception:
+                    pass
+            perf = build_performance_summary(
+                total_duration_ms=latency_ms,
+                execution_steps=[{
+                    "name": "LLM Call 1",
+                    "duration_ms": latency_ms,
+                    "start_time": start_t,
+                    "end_time": end_t,
+                }],
+                has_errors=error is not None,
+            )
             tracer._send(
                 name=name,
                 input=_safe_serialize(input_messages),
@@ -87,6 +108,9 @@ def _patch_create(
                 model=model,
                 metadata=metadata,
                 session_id=session_id,
+                performance_summary=perf,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
             )
 
     patched_create._agentx_patched = True
@@ -105,7 +129,7 @@ def _patch_stream(
         return
 
     def patched_stream(*args, **kwargs):
-        start = time.time()
+        start_t = time.time()
         ctx = original_stream(*args, **kwargs)
 
         class _TracedStream:
@@ -116,14 +140,31 @@ def _patch_stream(
 
             def __exit__(self_inner, exc_type, exc_val, tb):
                 result = ctx.__exit__(exc_type, exc_val, tb)
-                latency_ms = int((time.time() - start) * 1000)
+                end_t = time.time()
+                latency_ms = int((end_t - start_t) * 1000)
                 error = str(exc_val) if exc_val else None
                 output = None
+                input_tokens = None
+                output_tokens = None
                 try:
                     final = ctx.get_final_message()
                     output = final.content[0].text if final.content else None
+                    usage = getattr(final, "usage", None)
+                    if usage is not None:
+                        input_tokens = getattr(usage, "input_tokens", None)
+                        output_tokens = getattr(usage, "output_tokens", None)
                 except Exception:
                     pass
+                perf = build_performance_summary(
+                    total_duration_ms=latency_ms,
+                    execution_steps=[{
+                        "name": "LLM Call 1",
+                        "duration_ms": latency_ms,
+                        "start_time": start_t,
+                        "end_time": end_t,
+                    }],
+                    has_errors=error is not None,
+                )
                 tracer._send(
                     name=name,
                     input=_safe_serialize(kwargs.get("messages")),
@@ -134,6 +175,9 @@ def _patch_stream(
                     model=kwargs.get("model"),
                     metadata=metadata,
                     session_id=session_id,
+                    performance_summary=perf,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
                 )
                 return result
 

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx/integrations/google_adk.py

@@ -24,6 +24,7 @@ import time
 from typing import Any, Dict, List, Optional
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
+from agentx.integrations._perf import build_performance_summary
 
 try:
     from google.adk.plugins.base_plugin import BasePlugin
@@ -56,7 +57,7 @@ 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.
+    tool calls, latency, and a performance_summary via the ADK plugin callbacks.
 
     Register via the ``plugins`` list when constructing the ADK ``Runner``.
     """
@@ -80,6 +81,10 @@ class AgentXADKPlugin(BasePlugin):
         self._pending_inputs: Dict[str, str] = {}
         # id(tool_context) → start time float
         self._tool_starts: Dict[int, float] = {}
+        # invocation_id → stack of model call start times (FIFO)
+        # ADK creates new CallbackContext objects for before/after model callbacks,
+        # so we cannot use id(callback_context) as a key — use invocation_id instead.
+        self._model_starts: Dict[str, List[float]] = {}
 
     # ------------------------------------------------------------------
     # Run lifecycle
@@ -106,6 +111,10 @@ class AgentXADKPlugin(BasePlugin):
             "model": None,
             "tool_calls": [],
             "error": None,
+            "execution_steps": [],
+            "perf_tool_calls": [],
+            "input_tokens": 0,
+            "output_tokens": 0,
         }
 
     async def after_run_callback(self, *, invocation_context: Any) -> None:
@@ -114,6 +123,11 @@ class AgentXADKPlugin(BasePlugin):
         if state is None:
             return
         latency_ms = int((time.time() - state["start"]) * 1000)
+        perf = build_performance_summary(
+            total_duration_ms=latency_ms,
+            execution_steps=state["execution_steps"],
+            tool_call_steps=state["perf_tool_calls"],
+        )
         self._tracer._send(
             name=state["name"],
             input=state["input"],
@@ -124,10 +138,13 @@ class AgentXADKPlugin(BasePlugin):
             tool_calls=state["tool_calls"] or None,
             metadata=self._metadata,
             session_id=self._session_id,
+            performance_summary=perf,
+            input_tokens=state["input_tokens"] or None,
+            output_tokens=state["output_tokens"] or None,
         )
 
     # ------------------------------------------------------------------
-    # Model callbacks — capture model name and output
+    # Model callbacks — capture model name, output, and LLM step timing
     # ------------------------------------------------------------------
 
     async def before_model_callback(
@@ -135,22 +152,52 @@ class AgentXADKPlugin(BasePlugin):
     ) -> 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)
+        if state and not state["model"]:
+            model = getattr(llm_request, "model", None)
+            if model:
+                state["model"] = str(model)
+        # Push start time onto the per-invocation stack.
+        # ADK creates different CallbackContext objects for before vs after, so
+        # id(callback_context) cannot be used as a key across the two calls.
+        if inv_id not in self._model_starts:
+            self._model_starts[inv_id] = []
+        self._model_starts[inv_id].append(time.time())
 
     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)
+        # Pop the earliest queued start time (FIFO — model calls are sequential)
+        starts = self._model_starts.get(inv_id, [])
+        start_t = starts.pop(0) if starts else None
+        end_t = time.time()
+
         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
 
+        # Execution step
+        if start_t is not None:
+            steps = state["execution_steps"]
+            steps.append({
+                "name": f"LLM Call {len(steps) + 1}",
+                "duration_ms": (end_t - start_t) * 1000,
+                "start_time": start_t,
+                "end_time": end_t,
+            })
+
+        # Token counts
+        usage = getattr(llm_response, "usage_metadata", None)
+        if usage is not None:
+            state["input_tokens"] += int(getattr(usage, "prompt_token_count", None) or 0)
+            state["output_tokens"] += int(getattr(usage, "candidates_token_count", None) or 0)
+
     # ------------------------------------------------------------------
     # Tool callbacks
     # ------------------------------------------------------------------
@@ -172,15 +219,24 @@ class AgentXADKPlugin(BasePlugin):
         state = self._runs.get(inv_id)
         if state is None:
             return
-        start = self._tool_starts.pop(id(tool_context), None)
+        start_t = self._tool_starts.pop(id(tool_context), None)
+        end_t = time.time()
+        tool_name = getattr(tool, "name", "unknown")
         tool_call: Dict[str, Any] = {
-            "name": getattr(tool, "name", "unknown"),
+            "name": tool_name,
             "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))
+        if start_t is not None:
+            tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
         state["tool_calls"].append(tool_call)
+        if start_t is not None:
+            state["perf_tool_calls"].append({
+                "name": tool_name,
+                "duration_ms": (end_t - start_t) * 1000,
+                "start_time": start_t,
+                "end_time": end_t,
+            })
 
     async def on_tool_error_callback(
         self,
@@ -194,12 +250,21 @@ class AgentXADKPlugin(BasePlugin):
         state = self._runs.get(inv_id)
         if state is None:
             return
-        start = self._tool_starts.pop(id(tool_context), None)
+        start_t = self._tool_starts.pop(id(tool_context), None)
+        end_t = time.time()
+        tool_name = getattr(tool, "name", "unknown")
         tool_call: Dict[str, Any] = {
-            "name": getattr(tool, "name", "unknown"),
+            "name": tool_name,
             "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))
+        if start_t is not None:
+            tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
         state["tool_calls"].append(tool_call)
+        if start_t is not None:
+            state["perf_tool_calls"].append({
+                "name": tool_name,
+                "duration_ms": (end_t - start_t) * 1000,
+                "start_time": start_t,
+                "end_time": end_t,
+            })

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx/integrations/google_genai.py

@@ -19,6 +19,7 @@ import time
 from typing import Any, Dict, Optional
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
+from agentx.integrations._perf import build_performance_summary
 
 
 def patch_genai_client(
@@ -77,7 +78,7 @@ def _patch_generate_content(
         return
 
     def patched(*args, **kwargs):
-        start = time.time()
+        start_t = time.time()
         error: Optional[str] = None
         response = None
         try:
@@ -87,10 +88,28 @@ def _patch_generate_content(
             error = str(exc)
             raise
         finally:
-            latency_ms = int((time.time() - start) * 1000)
+            end_t = time.time()
+            latency_ms = int((end_t - start_t) * 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
+            input_tokens = None
+            output_tokens = None
+            if response is not None:
+                usage = getattr(response, "usage_metadata", None)
+                if usage is not None:
+                    input_tokens = getattr(usage, "prompt_token_count", None)
+                    output_tokens = getattr(usage, "candidates_token_count", None)
+            perf = build_performance_summary(
+                total_duration_ms=latency_ms,
+                execution_steps=[{
+                    "name": "LLM Call 1",
+                    "duration_ms": latency_ms,
+                    "start_time": start_t,
+                    "end_time": end_t,
+                }],
+                has_errors=error is not None,
+            )
             tracer._send(
                 name=name,
                 input=contents if isinstance(contents, str) else _safe_serialize(contents),
@@ -101,6 +120,9 @@ def _patch_generate_content(
                 model=str(model) if model else None,
                 metadata=metadata,
                 session_id=session_id,
+                performance_summary=perf,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
             )
 
     patched._agentx_patched = True
@@ -119,8 +141,9 @@ def _patch_generate_content_stream(
         return
 
     def patched_stream(*args, **kwargs):
-        start = time.time()
+        start_t = time.time()
         accumulated_text: list[str] = []
+        last_usage_metadata = None
         error: Optional[str] = None
 
         try:
@@ -128,14 +151,34 @@ def _patch_generate_content_stream(
                 text = getattr(chunk, "text", None)
                 if text:
                     accumulated_text.append(text)
+                # Track usage_metadata from the last chunk (Gemini includes it there)
+                chunk_usage = getattr(chunk, "usage_metadata", None)
+                if chunk_usage is not None:
+                    last_usage_metadata = chunk_usage
                 yield chunk
         except Exception as exc:
             error = str(exc)
             raise
         finally:
-            latency_ms = int((time.time() - start) * 1000)
+            end_t = time.time()
+            latency_ms = int((end_t - start_t) * 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)
+            input_tokens = None
+            output_tokens = None
+            if last_usage_metadata is not None:
+                input_tokens = getattr(last_usage_metadata, "prompt_token_count", None)
+                output_tokens = getattr(last_usage_metadata, "candidates_token_count", None)
+            perf = build_performance_summary(
+                total_duration_ms=latency_ms,
+                execution_steps=[{
+                    "name": "LLM Call 1",
+                    "duration_ms": latency_ms,
+                    "start_time": start_t,
+                    "end_time": end_t,
+                }],
+                has_errors=error is not None,
+            )
             tracer._send(
                 name=name,
                 input=contents if isinstance(contents, str) else _safe_serialize(contents),
@@ -146,6 +189,9 @@ def _patch_generate_content_stream(
                 model=str(model) if model else None,
                 metadata=metadata,
                 session_id=session_id,
+                performance_summary=perf,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
             )
 
     patched_stream._agentx_patched = True

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx/integrations/langchain.py

@@ -25,6 +25,7 @@ from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
+from agentx.integrations._perf import build_performance_summary
 
 try:
     from langchain_core.callbacks.base import BaseCallbackHandler
@@ -62,7 +63,6 @@ def _extract_output(outputs: Any) -> Any:
                 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:
@@ -177,6 +177,11 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         self._top_level: Dict[UUID, bool] = {}
         # Full parent-chain map so _find_top_ancestor can walk arbitrary depth
         self._parents: Dict[UUID, Optional[UUID]] = {}
+        # Retrieval steps that fire before on_chain_start (pre-run RAG pattern).
+        # Consumed and attached when the next top-level chain starts.
+        self._pending_retrieval_steps: List[Dict[str, Any]] = []
+        # run_id → {"start": float, "query": str}
+        self._retrieval_starts: Dict[UUID, Dict[str, Any]] = {}
 
     # ------------------------------------------------------------------
     # Chain lifecycle
@@ -195,11 +200,20 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         is_top = parent_run_id is None
         self._top_level[run_id] = is_top
         if is_top:
+            # Consume any retrieval steps that ran before this chain started
+            # (pre-run RAG: retriever.invoke() called before agent.invoke())
+            pending = self._pending_retrieval_steps[:]
+            self._pending_retrieval_steps.clear()
             self._runs[run_id] = {
                 "start": time.time(),
                 "input": _extract_input(inputs),
                 "tool_calls": [],
                 "model": None,
+                "execution_steps": [],
+                "perf_tool_calls": [],
+                "retrieval_steps": pending,
+                "input_tokens": 0,
+                "output_tokens": 0,
             }
 
     def on_chain_end(
@@ -217,8 +231,13 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             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)
+        perf = build_performance_summary(
+            total_duration_ms=latency_ms,
+            execution_steps=state["execution_steps"],
+            tool_call_steps=state["perf_tool_calls"],
+            retrieval_steps=state["retrieval_steps"],
+        )
         self._tracer._send(
             name=self._name,
             input=state["input"],
@@ -229,6 +248,9 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             tool_calls=tool_calls or None,
             metadata=self._metadata,
             session_id=self._session_id,
+            performance_summary=perf,
+            input_tokens=state["input_tokens"] or None,
+            output_tokens=state["output_tokens"] or None,
         )
         self._top_level.pop(run_id, None)
         self._parents.pop(run_id, None)
@@ -247,6 +269,13 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         if state is None:
             return
         latency_ms = int((time.time() - state["start"]) * 1000)
+        perf = build_performance_summary(
+            total_duration_ms=latency_ms,
+            execution_steps=state["execution_steps"],
+            tool_call_steps=state["perf_tool_calls"],
+            retrieval_steps=state["retrieval_steps"],
+            has_errors=True,
+        )
         self._tracer._send(
             name=self._name,
             input=state["input"],
@@ -257,6 +286,9 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             tool_calls=state["tool_calls"] or None,
             metadata=self._metadata,
             session_id=self._session_id,
+            performance_summary=perf,
+            input_tokens=state["input_tokens"] or None,
+            output_tokens=state["output_tokens"] or None,
         )
         self._top_level.pop(run_id, None)
         self._parents.pop(run_id, None)
@@ -265,6 +297,29 @@ class AgentXCallbackHandler(BaseCallbackHandler):
     # LLM lifecycle
     # ------------------------------------------------------------------
 
+    def _record_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        run_id: UUID,
+        parent_run_id: Optional[UUID],
+        kwargs: Dict[str, Any],
+    ) -> None:
+        """Shared logic for on_llm_start and on_chat_model_start."""
+        self._parents[run_id] = parent_run_id
+        kw = serialized.get("kwargs", {})
+        model = (
+            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")
+        )
+        model = str(model) if model and model not in ("None", "none") else None
+        self._runs[run_id] = {"llm_start": time.time(), "model": model}
+        top = self._find_top_ancestor(parent_run_id)
+        if top and not self._runs[top].get("model") and model:
+            self._runs[top]["model"] = model
+
     def on_llm_start(
         self,
         serialized: Dict[str, Any],
@@ -274,22 +329,18 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs,
     ) -> None:
-        self._parents[run_id] = parent_run_id
-        self._runs.setdefault(run_id, {})["llm_start"] = time.time()
+        self._record_llm_start(serialized, run_id, parent_run_id, kwargs)
 
-        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 = (
-                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 and model not in ("None", "none"):
-                self._runs[top]["model"] = str(model)
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ) -> None:
+        self._record_llm_start(serialized, run_id, parent_run_id, kwargs)
 
     def on_llm_end(
         self,
@@ -299,8 +350,40 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs,
     ) -> None:
-        self._runs.pop(run_id, None)
+        llm_state = self._runs.pop(run_id, None)
         self._parents.pop(run_id, None)
+        if llm_state:
+            start_t = llm_state.get("llm_start")
+            end_t = time.time()
+            top = self._find_top_ancestor(parent_run_id)
+            if start_t is not None and top and top in self._runs:
+                steps = self._runs[top]["execution_steps"]
+                steps.append({
+                    "name": f"LLM Call {len(steps) + 1}",
+                    "duration_ms": (end_t - start_t) * 1000,
+                    "start_time": start_t,
+                    "end_time": end_t,
+                })
+            # Extract token usage from LLMResult
+            if top and top in self._runs:
+                usage = {}
+                if hasattr(response, "llm_output") and isinstance(response.llm_output, dict):
+                    usage = response.llm_output.get("token_usage") or response.llm_output.get("usage") or {}
+                # Also check generations for token counts (some providers put it there)
+                if not usage and hasattr(response, "generations"):
+                    for gen_list in (response.generations or []):
+                        for gen in (gen_list or []):
+                            gen_info = getattr(gen, "generation_info", None) or {}
+                            if gen_info.get("prompt_tokens") or gen_info.get("completion_tokens"):
+                                usage = gen_info
+                                break
+                if usage:
+                    self._runs[top]["input_tokens"] += int(
+                        usage.get("prompt_tokens") or usage.get("input_tokens") or usage.get("prompt_token_count") or 0
+                    )
+                    self._runs[top]["output_tokens"] += int(
+                        usage.get("completion_tokens") or usage.get("output_tokens") or usage.get("candidates_token_count") or 0
+                    )
 
     # ------------------------------------------------------------------
     # Tool lifecycle
@@ -334,7 +417,9 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         self._parents.pop(run_id, None)
         if state is None:
             return
-        latency_ms = int((time.time() - state["start"]) * 1000)
+        end_t = time.time()
+        start_t = state["start"]
+        latency_ms = int((end_t - start_t) * 1000)
         tool_call = {
             "name": state["tool_name"],
             "input": state["tool_input"],
@@ -344,6 +429,12 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         top = self._find_top_ancestor(parent_run_id)
         if top and top in self._runs:
             self._runs[top]["tool_calls"].append(tool_call)
+            self._runs[top]["perf_tool_calls"].append({
+                "name": state["tool_name"],
+                "duration_ms": (end_t - start_t) * 1000,
+                "start_time": start_t,
+                "end_time": end_t,
+            })
 
     def on_tool_error(
         self,
@@ -357,15 +448,88 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         self._parents.pop(run_id, None)
         if state is None:
             return
+        end_t = time.time()
+        start_t = state.get("start", end_t)
         tool_call = {
             "name": state.get("tool_name", "unknown"),
             "input": state.get("tool_input"),
             "output": f"ERROR: {error}",
-            "latency_ms": int((time.time() - state["start"]) * 1000),
+            "latency_ms": int((end_t - start_t) * 1000),
         }
         top = self._find_top_ancestor(parent_run_id)
         if top and top in self._runs:
             self._runs[top]["tool_calls"].append(tool_call)
+            self._runs[top]["perf_tool_calls"].append({
+                "name": state.get("tool_name", "unknown"),
+                "duration_ms": (end_t - start_t) * 1000,
+                "start_time": start_t,
+                "end_time": end_t,
+            })
+
+    # ------------------------------------------------------------------
+    # Retriever lifecycle
+    # ------------------------------------------------------------------
+
+    def on_retriever_start(
+        self,
+        serialized: Dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ) -> None:
+        self._parents[run_id] = parent_run_id
+        self._retrieval_starts[run_id] = {"start": time.time(), "query": query}
+
+    def on_retriever_end(
+        self,
+        documents: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ) -> None:
+        state = self._retrieval_starts.pop(run_id, None)
+        self._parents.pop(run_id, None)
+        if state is None:
+            return
+        end_t = time.time()
+        start_t = state["start"]
+        query: Optional[str] = state["query"] or None
+        doc_count = len(documents) if hasattr(documents, "__len__") else None
+        step: Dict[str, Any] = {
+            "name": "Retrieval 1",  # renumbered below
+            "duration_ms": (end_t - start_t) * 1000,
+            "start_time": start_t,
+            "end_time": end_t,
+        }
+        if query:
+            step["query"] = query
+        if doc_count is not None:
+            step["doc_count"] = doc_count
+
+        top = self._find_top_ancestor(parent_run_id)
+        if top and top in self._runs:
+            # Retriever ran inside an active chain — attach directly
+            retrievals = self._runs[top]["retrieval_steps"]
+            step["name"] = f"Retrieval {len(retrievals) + 1}"
+            retrievals.append(step)
+        else:
+            # Retriever ran before the chain started (pre-run RAG pattern)
+            step["name"] = f"Retrieval {len(self._pending_retrieval_steps) + 1}"
+            self._pending_retrieval_steps.append(step)
+
+    def on_retriever_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs,
+    ) -> None:
+        self._retrieval_starts.pop(run_id, None)
+        self._parents.pop(run_id, None)
 
     # ------------------------------------------------------------------
     # Helpers

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx/integrations/openai_agents.py

@@ -20,6 +20,7 @@ from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
+from agentx.integrations._perf import build_performance_summary
 
 
 def _iso_to_ts(iso: Optional[str]) -> Optional[float]:
@@ -143,6 +144,10 @@ class AgentXTracingProcessor:
             "output": None,
             "tool_calls": [],
             "model": None,
+            "execution_steps": [],
+            "perf_tool_calls": [],
+            "input_tokens": 0,
+            "output_tokens": 0,
         }
 
     def on_trace_end(self, trace: Any) -> None:
@@ -151,6 +156,11 @@ class AgentXTracingProcessor:
         if state is None:
             return
         latency_ms = int((time.time() - state["start"]) * 1000)
+        perf = build_performance_summary(
+            total_duration_ms=latency_ms,
+            execution_steps=state["execution_steps"],
+            tool_call_steps=state["perf_tool_calls"],
+        )
         self._tracer._send(
             name=state["name"],
             input=state.get("input"),
@@ -161,6 +171,9 @@ class AgentXTracingProcessor:
             tool_calls=state["tool_calls"] or None,
             metadata=self._metadata,
             session_id=self._session_id,
+            performance_summary=perf,
+            input_tokens=state["input_tokens"] or None,
+            output_tokens=state["output_tokens"] or None,
         )
 
     def on_span_start(self, span: Any) -> None:
@@ -178,6 +191,9 @@ class AgentXTracingProcessor:
         state = self._spans[trace_id]
         span_type = getattr(span_data, "type", None)
 
+        t0 = _iso_to_ts(getattr(span, "started_at", None))
+        t1 = _iso_to_ts(getattr(span, "ended_at", None))
+
         if span_type == "generation":
             # Capture input from the first generation span
             if state["input"] is None and span_data.input:
@@ -188,6 +204,20 @@ class AgentXTracingProcessor:
             # Capture model name
             if not state["model"] and span_data.model:
                 state["model"] = str(span_data.model)
+            # Execution step
+            if t0 is not None and t1 is not None:
+                steps = state["execution_steps"]
+                steps.append({
+                    "name": f"LLM Call {len(steps) + 1}",
+                    "duration_ms": (t1 - t0) * 1000,
+                    "start_time": t0,
+                    "end_time": t1,
+                })
+            # Token counts — usage is a dict with "input_tokens" / "output_tokens"
+            usage = getattr(span_data, "usage", None)
+            if isinstance(usage, dict):
+                state["input_tokens"] += int(usage.get("input_tokens") or 0)
+                state["output_tokens"] += int(usage.get("output_tokens") or 0)
 
         elif span_type == "response":
             # Responses API path — extract from the response object
@@ -204,18 +234,43 @@ class AgentXTracingProcessor:
                     model = getattr(response, "model", None)
                     if model:
                         state["model"] = str(model)
+                # Token counts from response.usage or span_data.usage
+                usage = getattr(response, "usage", None) or getattr(span_data, "usage", None)
+                if isinstance(usage, dict):
+                    state["input_tokens"] += int(usage.get("input_tokens") or 0)
+                    state["output_tokens"] += int(usage.get("output_tokens") or 0)
+                elif usage is not None:
+                    state["input_tokens"] += int(getattr(usage, "input_tokens", None) or 0)
+                    state["output_tokens"] += int(getattr(usage, "output_tokens", None) or 0)
+            # Execution step
+            if t0 is not None and t1 is not None:
+                steps = state["execution_steps"]
+                steps.append({
+                    "name": f"LLM Call {len(steps) + 1}",
+                    "duration_ms": (t1 - t0) * 1000,
+                    "start_time": t0,
+                    "end_time": t1,
+                })
 
         elif span_type == "function":
             # Tool / function call
+            latency = _span_latency_ms(span)
             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)
+            # Perf tool call with timestamps
+            if t0 is not None and t1 is not None:
+                state["perf_tool_calls"].append({
+                    "name": span_data.name,
+                    "duration_ms": (t1 - t0) * 1000,
+                    "start_time": t0,
+                    "end_time": t1,
+                })
 
     def force_flush(self) -> None:
         self._tracer.flush()

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx/tracing/tracer.py

@@ -143,8 +143,9 @@ class _TraceSpan:
     def _wrap_sync(self, fn: F) -> F:
         @functools.wraps(fn)
         def wrapper(*args, **kwargs):
+            from agentx.integrations._perf import build_performance_summary
             captured_input = _capture_fn_input(fn, args, kwargs)
-            start = time.time()
+            start_t = time.time()
             error: Optional[str] = None
             output = None
             try:
@@ -154,16 +155,28 @@ class _TraceSpan:
                 error = str(exc)
                 raise
             finally:
+                end_t = time.time()
+                latency_ms = int((end_t - start_t) * 1000)
                 self._tracer._send(
                     name=self.name,
                     input=captured_input,
                     output=_safe_serialize(output) if output is not None else None,
-                    latency_ms=int((time.time() - start) * 1000),
+                    latency_ms=latency_ms,
                     error=error,
                     metadata=self._metadata,
                     framework=self._framework,
                     model=self._model,
                     session_id=self._session_id,
+                    performance_summary=build_performance_summary(
+                        total_duration_ms=latency_ms,
+                        execution_steps=[{
+                            "name": "LLM Call 1",
+                            "duration_ms": latency_ms,
+                            "start_time": start_t,
+                            "end_time": end_t,
+                        }],
+                        has_errors=error is not None,
+                    ),
                 )
 
         return wrapper  # type: ignore[return-value]
@@ -171,8 +184,9 @@ class _TraceSpan:
     def _wrap_async(self, fn: F) -> F:
         @functools.wraps(fn)
         async def wrapper(*args, **kwargs):
+            from agentx.integrations._perf import build_performance_summary
             captured_input = _capture_fn_input(fn, args, kwargs)
-            start = time.time()
+            start_t = time.time()
             error: Optional[str] = None
             output = None
             try:
@@ -182,16 +196,28 @@ class _TraceSpan:
                 error = str(exc)
                 raise
             finally:
+                end_t = time.time()
+                latency_ms = int((end_t - start_t) * 1000)
                 self._tracer._send(
                     name=self.name,
                     input=captured_input,
                     output=_safe_serialize(output) if output is not None else None,
-                    latency_ms=int((time.time() - start) * 1000),
+                    latency_ms=latency_ms,
                     error=error,
                     metadata=self._metadata,
                     framework=self._framework,
                     model=self._model,
                     session_id=self._session_id,
+                    performance_summary=build_performance_summary(
+                        total_duration_ms=latency_ms,
+                        execution_steps=[{
+                            "name": "LLM Call 1",
+                            "duration_ms": latency_ms,
+                            "start_time": start_t,
+                            "end_time": end_t,
+                        }],
+                        has_errors=error is not None,
+                    ),
                 )
 
         return wrapper  # type: ignore[return-value]
@@ -462,4 +488,10 @@ class Tracer:
             wire["metadata"] = payload["metadata"]
         if "session_id" in payload:
             wire["session_id"] = payload["session_id"]
+        if "performance_summary" in payload:
+            wire["performance_summary"] = payload["performance_summary"]
+        if "input_tokens" in payload:
+            wire["input_tokens"] = payload["input_tokens"]
+        if "output_tokens" in payload:
+            wire["output_tokens"] = payload["output_tokens"]
         self._client.enqueue(wire)

agentx_python-0.6.1/agentx/version.py

@@ -0,0 +1 @@
+VERSION = "0.6.1"

{agentx_python-0.5.2 → agentx_python-0.6.1/agentx_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.5.2
+Version: 0.6.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

{agentx_python-0.5.2 → agentx_python-0.6.1}/agentx_python.egg-info/SOURCES.txt

@@ -21,6 +21,7 @@ agentx/evaluations/adapters/http_endpoint.py
 agentx/evaluations/adapters/precomputed.py
 agentx/evaluations/adapters/raw.py
 agentx/integrations/__init__.py
+agentx/integrations/_perf.py
 agentx/integrations/anthropic.py
 agentx/integrations/crewai.py
 agentx/integrations/google_adk.py

agentx_python-0.5.2/agentx/version.py

@@ -1 +0,0 @@
-VERSION = "0.5.2"