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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {agentx_python-0.5.2/agentx_python.egg-info → agentx_python-0.6.2}/PKG-INFO +1 -1
  2. agentx_python-0.6.2/agentx/integrations/_perf.py +226 -0
  3. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/integrations/anthropic.py +89 -16
  4. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/integrations/crewai.py +31 -2
  5. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/integrations/google_adk.py +113 -16
  6. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/integrations/google_genai.py +66 -7
  7. agentx_python-0.6.2/agentx/integrations/langchain.py +659 -0
  8. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/integrations/openai_agents.py +106 -17
  9. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/tracing/tracer.py +369 -6
  10. agentx_python-0.6.2/agentx/version.py +1 -0
  11. {agentx_python-0.5.2 → agentx_python-0.6.2/agentx_python.egg-info}/PKG-INFO +1 -1
  12. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx_python.egg-info/SOURCES.txt +1 -0
  13. agentx_python-0.5.2/agentx/integrations/langchain.py +0 -383
  14. agentx_python-0.5.2/agentx/version.py +0 -1
  15. {agentx_python-0.5.2 → agentx_python-0.6.2}/LICENSE +0 -0
  16. {agentx_python-0.5.2 → agentx_python-0.6.2}/README.md +0 -0
  17. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/__init__.py +0 -0
  18. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/agentx.py +0 -0
  19. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/__init__.py +0 -0
  20. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/_term.py +0 -0
  21. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/adapters/__init__.py +0 -0
  22. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/adapters/http_endpoint.py +0 -0
  23. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/adapters/precomputed.py +0 -0
  24. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/adapters/raw.py +0 -0
  25. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/client.py +0 -0
  26. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/datasets.py +0 -0
  27. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/models.py +0 -0
  28. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/redaction.py +0 -0
  29. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/reporting.py +0 -0
  30. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/results.py +0 -0
  31. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/runner.py +0 -0
  32. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/evaluations/tracing.py +0 -0
  33. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/exceptions.py +0 -0
  34. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/integrations/__init__.py +0 -0
  35. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/resources/__init__.py +0 -0
  36. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/resources/agent.py +0 -0
  37. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/resources/conversation.py +0 -0
  38. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/resources/workforce.py +0 -0
  39. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/tracing/__init__.py +0 -0
  40. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/tracing/ci_types.py +0 -0
  41. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/tracing/ingest_client.py +0 -0
  42. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx/util.py +0 -0
  43. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx_python.egg-info/dependency_links.txt +0 -0
  44. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx_python.egg-info/requires.txt +0 -0
  45. {agentx_python-0.5.2 → agentx_python-0.6.2}/agentx_python.egg-info/top_level.txt +0 -0
  46. {agentx_python-0.5.2 → agentx_python-0.6.2}/setup.cfg +0 -0
  47. {agentx_python-0.5.2 → agentx_python-0.6.2}/setup.py +0 -0
  48. {agentx_python-0.5.2 → agentx_python-0.6.2}/tests/test_integration.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentx-python
3
- Version: 0.5.2
3
+ Version: 0.6.2
4
4
  Summary: Official Python SDK for AgentX (https://www.agentx.so/)
5
5
  Home-page: https://github.com/AgentX-ai/AgentX-python
6
6
  Author: Robin Wang and AgentX Team
@@ -0,0 +1,226 @@
1
+ """
2
+ Shared helper for building a performance_summary that matches the AgentX
3
+ platform agent format.
4
+
5
+ External framework integrations collect:
6
+ - execution_steps — LLM inference calls (name, duration_ms, start_time, end_time)
7
+ - tool_call_steps — function/tool invocations (same shape)
8
+
9
+ ``build_performance_summary`` merges them into the full structure expected by
10
+ the backend ingest endpoint and the AgentX UI.
11
+ """
12
+ from __future__ import annotations
13
+
14
+ from typing import Any, Dict, List, Optional
15
+
16
+
17
+ def build_performance_summary(
18
+ total_duration_ms: float,
19
+ execution_steps: Optional[List[Dict[str, Any]]] = None,
20
+ tool_call_steps: Optional[List[Dict[str, Any]]] = None,
21
+ retrieval_steps: Optional[List[Dict[str, Any]]] = None,
22
+ has_errors: bool = False,
23
+ ) -> Dict[str, Any]:
24
+ """
25
+ Return a ``performance_summary`` dict compatible with the AgentX platform
26
+ agent format.
27
+
28
+ Args:
29
+ total_duration_ms: Wall-clock duration of the entire run.
30
+ execution_steps: LLM call entries — each must have ``name``,
31
+ ``duration_ms``, and optionally ``start_time`` /
32
+ ``end_time`` as Unix floats, plus ``model``,
33
+ ``input``, ``output``, ``inputTokenSize``,
34
+ ``outputTokenSize`` for per-call drill-down.
35
+ tool_call_steps: Tool / function call entries — same shape, plus
36
+ optional ``input`` / ``output``.
37
+ retrieval_steps: RAG / vector-store retrieval entries — same shape,
38
+ plus optional ``query`` (str), ``doc_count`` (int)
39
+ and ``output`` (the retrieved content).
40
+ has_errors: Set to True when the run ended with an exception.
41
+
42
+ None of these fields are truncated here — callers decide what to capture.
43
+ """
44
+ execution_steps = execution_steps or []
45
+ tool_call_steps = tool_call_steps or []
46
+ retrieval_steps = retrieval_steps or []
47
+
48
+ # Merge all items with their phase type, sort by start_time
49
+ all_items: List[tuple[Dict[str, Any], str]] = (
50
+ [(s, "execution_step") for s in execution_steps]
51
+ + [(t, "tool_call") for t in tool_call_steps]
52
+ + [(r, "retrieval") for r in retrieval_steps]
53
+ )
54
+ all_items.sort(key=lambda x: x[0].get("start_time") or 0)
55
+
56
+ unified: List[Dict[str, Any]] = []
57
+ out_steps: List[Dict[str, Any]] = []
58
+ out_tools: List[Dict[str, Any]] = []
59
+ out_retrievals: List[Dict[str, Any]] = []
60
+
61
+ for order, (item, phase_type) in enumerate(all_items, start=1):
62
+ entry: Dict[str, Any] = {
63
+ "name": item["name"],
64
+ "duration_ms": round(float(item["duration_ms"]), 3),
65
+ "start_order": order,
66
+ "phase_type": phase_type,
67
+ }
68
+ if item.get("start_time") is not None:
69
+ entry["start_time"] = item["start_time"]
70
+ if item.get("end_time") is not None:
71
+ entry["end_time"] = item["end_time"]
72
+ # Per-step drill-down detail — untruncated, present when the caller
73
+ # captured it (e.g. an LLM call's model/prompt/completion, a tool
74
+ # call's args/result, or a retrieval's returned documents).
75
+ if item.get("model"):
76
+ entry["model"] = item["model"]
77
+ if item.get("input") is not None:
78
+ entry["input"] = item["input"]
79
+ if item.get("output") is not None:
80
+ entry["output"] = item["output"]
81
+ # Per-call token usage — populated for LLM call execution_steps by
82
+ # integrations that read it off the provider response (e.g. Anthropic's
83
+ # response.usage), so performance_summary carries a token breakdown
84
+ # per call, not just the trace-level input_tokens/output_tokens total.
85
+ # Named to match PromptTrace.inputTokenSize/outputTokenSize on the backend.
86
+ if item.get("inputTokenSize") is not None:
87
+ entry["inputTokenSize"] = item["inputTokenSize"]
88
+ if item.get("outputTokenSize") is not None:
89
+ entry["outputTokenSize"] = item["outputTokenSize"]
90
+ # Retrieval-specific fields
91
+ if phase_type == "retrieval":
92
+ if item.get("query"):
93
+ entry["query"] = item["query"]
94
+ if item.get("doc_count") is not None:
95
+ entry["doc_count"] = item["doc_count"]
96
+ unified.append(entry)
97
+
98
+ flat = {k: v for k, v in entry.items() if k != "phase_type"}
99
+ if phase_type == "execution_step":
100
+ out_steps.append(flat)
101
+ elif phase_type == "tool_call":
102
+ out_tools.append(flat)
103
+ else:
104
+ out_retrievals.append(flat)
105
+
106
+ tools_total_ms = round(sum(t["duration_ms"] for t in out_tools), 3)
107
+ retrievals_total_ms = round(sum(r["duration_ms"] for r in out_retrievals), 3)
108
+
109
+ return {
110
+ "total_duration_ms": round(float(total_duration_ms), 3),
111
+ "todo_tasks_enabled": None,
112
+ "main_phases": [],
113
+ "tool_calls": out_tools,
114
+ "mcp_tool_calls": [],
115
+ "action_tool_calls": [],
116
+ "delegate_calls": [],
117
+ "execution_steps": out_steps,
118
+ "knowledge_retrievals": out_retrievals,
119
+ "todo_operations": [],
120
+ "detailed_phases": [],
121
+ "active_phases": [],
122
+ "has_errors": has_errors,
123
+ "memory_actions": None,
124
+ "unified_timeline": unified,
125
+ "statistics": {
126
+ "total_main_phases": 0,
127
+ "total_tool_calls": len(out_tools),
128
+ "total_mcp_tool_calls": 0,
129
+ "total_action_tool_calls": 0,
130
+ "total_delegate_calls": 0,
131
+ "total_execution_steps": len(out_steps),
132
+ "total_knowledge_retrievals": len(out_retrievals),
133
+ "total_todo_operations": 0,
134
+ "main_phases_total_ms": 0,
135
+ "tool_calls_total_ms": tools_total_ms,
136
+ "knowledge_retrievals_total_ms": retrievals_total_ms,
137
+ "mcp_tool_calls_total_ms": 0,
138
+ "todo_operations_total_ms": 0,
139
+ },
140
+ }
141
+
142
+
143
+ def _merge_steps(
144
+ summary: Dict[str, Any],
145
+ steps: List[Dict[str, Any]],
146
+ *,
147
+ phase_type: str,
148
+ target_key: str,
149
+ stats_count_key: str,
150
+ stats_total_ms_key: str,
151
+ extra_fields: tuple = (),
152
+ default_name_prefix: str = "Step",
153
+ ) -> Dict[str, Any]:
154
+ """Fold externally-recorded steps into an already-built ``performance_summary``."""
155
+ if not steps:
156
+ return summary
157
+
158
+ next_order = len(summary.get("unified_timeline", [])) + 1
159
+ for i, item in enumerate(steps):
160
+ entry: Dict[str, Any] = {
161
+ "name": item.get("name") or f"{default_name_prefix} {i + 1}",
162
+ "duration_ms": round(float(item.get("duration_ms") or 0), 3),
163
+ "start_order": next_order + i,
164
+ "phase_type": phase_type,
165
+ }
166
+ if item.get("start_time") is not None:
167
+ entry["start_time"] = item["start_time"]
168
+ if item.get("end_time") is not None:
169
+ entry["end_time"] = item["end_time"]
170
+ for field in extra_fields:
171
+ value = item.get(field)
172
+ if value not in (None, ""):
173
+ entry[field] = value
174
+
175
+ summary["unified_timeline"].append(entry)
176
+ summary[target_key].append({k: v for k, v in entry.items() if k != "phase_type"})
177
+
178
+ summary["unified_timeline"].sort(key=lambda x: x.get("start_time") or 0)
179
+ summary["statistics"][stats_count_key] = len(summary[target_key])
180
+ summary["statistics"][stats_total_ms_key] = round(
181
+ sum(r["duration_ms"] for r in summary[target_key]), 3
182
+ )
183
+ return summary
184
+
185
+
186
+ def merge_retrieval_steps(
187
+ summary: Dict[str, Any],
188
+ retrieval_steps: List[Dict[str, Any]],
189
+ ) -> Dict[str, Any]:
190
+ """
191
+ Fold externally-recorded retrieval steps (e.g. from ``Tracer.record_retrieval``,
192
+ used for hand-rolled RAG lookups the framework integration can't see) into an
193
+ already-built ``performance_summary``.
194
+ """
195
+ return _merge_steps(
196
+ summary,
197
+ retrieval_steps,
198
+ phase_type="retrieval",
199
+ target_key="knowledge_retrievals",
200
+ stats_count_key="total_knowledge_retrievals",
201
+ stats_total_ms_key="knowledge_retrievals_total_ms",
202
+ extra_fields=("query", "doc_count", "output"),
203
+ default_name_prefix="Retrieval",
204
+ )
205
+
206
+
207
+ def merge_tool_call_steps(
208
+ summary: Dict[str, Any],
209
+ tool_call_steps: List[Dict[str, Any]],
210
+ ) -> Dict[str, Any]:
211
+ """
212
+ Fold externally-recorded tool calls (e.g. from ``Tracer.record_tool_call``,
213
+ used when a tool executes outside a framework integration's visibility —
214
+ such as a manual Anthropic tool-use loop) into an already-built
215
+ ``performance_summary``.
216
+ """
217
+ return _merge_steps(
218
+ summary,
219
+ tool_call_steps,
220
+ phase_type="tool_call",
221
+ target_key="tool_calls",
222
+ stats_count_key="total_tool_calls",
223
+ stats_total_ms_key="tool_calls_total_ms",
224
+ extra_fields=("input", "output"),
225
+ default_name_prefix="Tool Call",
226
+ )
@@ -19,6 +19,7 @@ import time
19
19
  from typing import Any, Dict, Optional
20
20
 
21
21
  from agentx.tracing.tracer import Tracer, _safe_serialize
22
+ from agentx.integrations._perf import build_performance_summary
22
23
 
23
24
 
24
25
  def patch_anthropic_client(
@@ -58,7 +59,7 @@ def _patch_create(
58
59
  return # already patched
59
60
 
60
61
  def patched_create(*args, **kwargs):
61
- start = time.time()
62
+ start_t = time.time()
62
63
  error: Optional[str] = None
63
64
  response = None
64
65
  try:
@@ -68,26 +69,73 @@ def _patch_create(
68
69
  error = str(exc)
69
70
  raise
70
71
  finally:
71
- latency_ms = int((time.time() - start) * 1000)
72
+ end_t = time.time()
73
+ latency_ms = int((end_t - start_t) * 1000)
72
74
  input_messages = kwargs.get("messages") or (args[0] if args else None)
73
75
  model = kwargs.get("model")
74
76
  output = None
77
+ input_tokens = None
78
+ output_tokens = None
75
79
  if response is not None:
76
80
  try:
77
81
  output = response.content[0].text if response.content else None
78
82
  except Exception:
79
- output = str(response)[:500]
80
- tracer._send(
81
- name=name,
82
- input=_safe_serialize(input_messages),
83
- output=output,
84
- latency_ms=latency_ms,
85
- error=error,
86
- framework="anthropic",
87
- model=model,
88
- metadata=metadata,
89
- session_id=session_id,
90
- )
83
+ output = str(response)
84
+ try:
85
+ usage = getattr(response, "usage", None)
86
+ if usage is not None:
87
+ input_tokens = getattr(usage, "input_tokens", None)
88
+ output_tokens = getattr(usage, "output_tokens", None)
89
+ except Exception:
90
+ pass
91
+
92
+ active_span = tracer.current_span
93
+ if active_span is not None:
94
+ # Part of a `with tracer.trace(...)` block (e.g. a multi-call
95
+ # agentic loop) — attach as one LLM-call step on that span's
96
+ # trace instead of sending an independent trace per call.
97
+ if error is not None:
98
+ active_span.set_error(error)
99
+ active_span._record_llm_call(
100
+ duration_ms=latency_ms,
101
+ start_time=start_t,
102
+ end_time=end_t,
103
+ input=_safe_serialize(input_messages),
104
+ output=output,
105
+ model=model,
106
+ input_tokens=input_tokens,
107
+ output_tokens=output_tokens,
108
+ )
109
+ else:
110
+ perf = build_performance_summary(
111
+ total_duration_ms=latency_ms,
112
+ execution_steps=[{
113
+ "name": "LLM Call 1",
114
+ "duration_ms": latency_ms,
115
+ "start_time": start_t,
116
+ "end_time": end_t,
117
+ "model": model,
118
+ "input": _safe_serialize(input_messages),
119
+ "output": output,
120
+ "inputTokenSize": input_tokens,
121
+ "outputTokenSize": output_tokens,
122
+ }],
123
+ has_errors=error is not None,
124
+ )
125
+ tracer._send(
126
+ name=name,
127
+ input=_safe_serialize(input_messages),
128
+ output=output,
129
+ latency_ms=latency_ms,
130
+ error=error,
131
+ framework="anthropic",
132
+ model=model,
133
+ metadata=metadata,
134
+ session_id=session_id,
135
+ performance_summary=perf,
136
+ input_tokens=input_tokens,
137
+ output_tokens=output_tokens,
138
+ )
91
139
 
92
140
  patched_create._agentx_patched = True
93
141
  messages_resource.create = patched_create
@@ -105,7 +153,7 @@ def _patch_stream(
105
153
  return
106
154
 
107
155
  def patched_stream(*args, **kwargs):
108
- start = time.time()
156
+ start_t = time.time()
109
157
  ctx = original_stream(*args, **kwargs)
110
158
 
111
159
  class _TracedStream:
@@ -116,14 +164,36 @@ def _patch_stream(
116
164
 
117
165
  def __exit__(self_inner, exc_type, exc_val, tb):
118
166
  result = ctx.__exit__(exc_type, exc_val, tb)
119
- latency_ms = int((time.time() - start) * 1000)
167
+ end_t = time.time()
168
+ latency_ms = int((end_t - start_t) * 1000)
120
169
  error = str(exc_val) if exc_val else None
121
170
  output = None
171
+ input_tokens = None
172
+ output_tokens = None
122
173
  try:
123
174
  final = ctx.get_final_message()
124
175
  output = final.content[0].text if final.content else None
176
+ usage = getattr(final, "usage", None)
177
+ if usage is not None:
178
+ input_tokens = getattr(usage, "input_tokens", None)
179
+ output_tokens = getattr(usage, "output_tokens", None)
125
180
  except Exception:
126
181
  pass
182
+ perf = build_performance_summary(
183
+ total_duration_ms=latency_ms,
184
+ execution_steps=[{
185
+ "name": "LLM Call 1",
186
+ "duration_ms": latency_ms,
187
+ "start_time": start_t,
188
+ "end_time": end_t,
189
+ "model": kwargs.get("model"),
190
+ "input": _safe_serialize(kwargs.get("messages")),
191
+ "output": output,
192
+ "inputTokenSize": input_tokens,
193
+ "outputTokenSize": output_tokens,
194
+ }],
195
+ has_errors=error is not None,
196
+ )
127
197
  tracer._send(
128
198
  name=name,
129
199
  input=_safe_serialize(kwargs.get("messages")),
@@ -134,6 +204,9 @@ def _patch_stream(
134
204
  model=kwargs.get("model"),
135
205
  metadata=metadata,
136
206
  session_id=session_id,
207
+ performance_summary=perf,
208
+ input_tokens=input_tokens,
209
+ output_tokens=output_tokens,
137
210
  )
138
211
  return result
139
212
 
@@ -22,6 +22,7 @@ import time
22
22
  from typing import Any, Dict, Optional
23
23
 
24
24
  from agentx.tracing.tracer import Tracer, _safe_serialize
25
+ from agentx.integrations._perf import build_performance_summary
25
26
 
26
27
 
27
28
  class AgentXCrewObserver:
@@ -64,15 +65,38 @@ class AgentXCrewObserver:
64
65
 
65
66
  # Collect task outputs as tool_calls for observability
66
67
  tool_calls = []
68
+ execution_steps = []
67
69
  if result is not None:
68
70
  task_outputs = getattr(result, "tasks_output", []) or []
69
71
  for task_out in task_outputs:
72
+ description = getattr(task_out, "description", "task")
73
+ name = description[:100] # display label only — full text goes in "input"
74
+ task_output = str(getattr(task_out, "raw", ""))
70
75
  tool_calls.append(
71
76
  {
72
- "name": getattr(task_out, "description", "task")[:100],
73
- "output": str(getattr(task_out, "raw", ""))[:500],
77
+ "name": name,
78
+ "input": description,
79
+ "output": task_output,
74
80
  }
75
81
  )
82
+ execution_steps.append({
83
+ "name": name,
84
+ "duration_ms": 0,
85
+ "input": description,
86
+ "output": task_output,
87
+ })
88
+
89
+ if not execution_steps:
90
+ # No per-task breakdown available — record the whole kickoff
91
+ # as a single step so the trace still gets timing detail.
92
+ execution_steps.append({"name": self._name, "duration_ms": latency_ms})
93
+ else:
94
+ # Task-level timing isn't exposed by CrewOutput; attribute the
95
+ # total latency evenly across tasks so the timeline still sums
96
+ # to the measured wall-clock duration.
97
+ per_step_ms = latency_ms / len(execution_steps)
98
+ for step in execution_steps:
99
+ step["duration_ms"] = per_step_ms
76
100
 
77
101
  self._tracer._send(
78
102
  name=self._name,
@@ -84,6 +108,11 @@ class AgentXCrewObserver:
84
108
  tool_calls=tool_calls or None,
85
109
  metadata=self._metadata,
86
110
  session_id=self._session_id,
111
+ performance_summary=build_performance_summary(
112
+ total_duration_ms=latency_ms,
113
+ execution_steps=execution_steps,
114
+ has_errors=error is not None,
115
+ ),
87
116
  )
88
117
 
89
118
  def observe(
@@ -24,6 +24,7 @@ import time
24
24
  from typing import Any, Dict, List, Optional
25
25
 
26
26
  from agentx.tracing.tracer import Tracer, _safe_serialize
27
+ from agentx.integrations._perf import build_performance_summary
27
28
 
28
29
  try:
29
30
  from google.adk.plugins.base_plugin import BasePlugin
@@ -51,12 +52,22 @@ def _content_to_text(content: Any) -> Optional[str]:
51
52
  return " ".join(texts) if texts else None
52
53
 
53
54
 
55
+ def _contents_to_text(contents: Any) -> Optional[str]:
56
+ """Extract plain text from an LlmRequest's ``contents`` (a list of Content)."""
57
+ if contents is None:
58
+ return None
59
+ if not isinstance(contents, (list, tuple)):
60
+ contents = [contents]
61
+ texts = [t for t in (_content_to_text(c) for c in contents) if t]
62
+ return "\n".join(texts) if texts else None
63
+
64
+
54
65
  class AgentXADKPlugin(BasePlugin):
55
66
  """
56
67
  Google ADK plugin that sends one AgentX trace per runner invocation.
57
68
 
58
69
  Captures input (user message), output (final model reply), model name,
59
- tool calls, and latency via the ADK plugin callback hooks.
70
+ tool calls, latency, and a performance_summary via the ADK plugin callbacks.
60
71
 
61
72
  Register via the ``plugins`` list when constructing the ADK ``Runner``.
62
73
  """
@@ -80,6 +91,10 @@ class AgentXADKPlugin(BasePlugin):
80
91
  self._pending_inputs: Dict[str, str] = {}
81
92
  # id(tool_context) → start time float
82
93
  self._tool_starts: Dict[int, float] = {}
94
+ # invocation_id → stack of model call start times (FIFO)
95
+ # ADK creates new CallbackContext objects for before/after model callbacks,
96
+ # so we cannot use id(callback_context) as a key — use invocation_id instead.
97
+ self._model_starts: Dict[str, List[float]] = {}
83
98
 
84
99
  # ------------------------------------------------------------------
85
100
  # Run lifecycle
@@ -106,6 +121,10 @@ class AgentXADKPlugin(BasePlugin):
106
121
  "model": None,
107
122
  "tool_calls": [],
108
123
  "error": None,
124
+ "execution_steps": [],
125
+ "perf_tool_calls": [],
126
+ "input_tokens": 0,
127
+ "output_tokens": 0,
109
128
  }
110
129
 
111
130
  async def after_run_callback(self, *, invocation_context: Any) -> None:
@@ -114,6 +133,11 @@ class AgentXADKPlugin(BasePlugin):
114
133
  if state is None:
115
134
  return
116
135
  latency_ms = int((time.time() - state["start"]) * 1000)
136
+ perf = build_performance_summary(
137
+ total_duration_ms=latency_ms,
138
+ execution_steps=state["execution_steps"],
139
+ tool_call_steps=state["perf_tool_calls"],
140
+ )
117
141
  self._tracer._send(
118
142
  name=state["name"],
119
143
  input=state["input"],
@@ -124,10 +148,13 @@ class AgentXADKPlugin(BasePlugin):
124
148
  tool_calls=state["tool_calls"] or None,
125
149
  metadata=self._metadata,
126
150
  session_id=self._session_id,
151
+ performance_summary=perf,
152
+ input_tokens=state["input_tokens"] or None,
153
+ output_tokens=state["output_tokens"] or None,
127
154
  )
128
155
 
129
156
  # ------------------------------------------------------------------
130
- # Model callbacks — capture model name and output
157
+ # Model callbacks — capture model name, output, and LLM step timing
131
158
  # ------------------------------------------------------------------
132
159
 
133
160
  async def before_model_callback(
@@ -135,22 +162,66 @@ class AgentXADKPlugin(BasePlugin):
135
162
  ) -> None:
136
163
  inv_id = callback_context.get_invocation_context().invocation_id
137
164
  state = self._runs.get(inv_id)
138
- if state and not state["model"] and llm_request.model:
139
- state["model"] = str(llm_request.model)
165
+ model = getattr(llm_request, "model", None)
166
+ model_str = str(model) if model else None
167
+ if state and not state["model"] and model_str:
168
+ state["model"] = model_str
169
+ # Push start time + this call's model/input onto the per-invocation
170
+ # stack, so after_model_callback can pair them back up.
171
+ # ADK creates different CallbackContext objects for before vs after, so
172
+ # id(callback_context) cannot be used as a key across the two calls.
173
+ if inv_id not in self._model_starts:
174
+ self._model_starts[inv_id] = []
175
+ self._model_starts[inv_id].append({
176
+ "start": time.time(),
177
+ "model": model_str,
178
+ "input": _contents_to_text(getattr(llm_request, "contents", None)),
179
+ })
140
180
 
141
181
  async def after_model_callback(
142
182
  self, *, callback_context: Any, llm_response: Any
143
183
  ) -> None:
144
184
  inv_id = callback_context.get_invocation_context().invocation_id
145
185
  state = self._runs.get(inv_id)
186
+ # Pop the earliest queued call (FIFO — model calls are sequential)
187
+ starts = self._model_starts.get(inv_id, [])
188
+ call_start = starts.pop(0) if starts else None
189
+ start_t = call_start.get("start") if call_start else None
190
+ end_t = time.time()
191
+
146
192
  if state is None:
147
193
  return
194
+
148
195
  content = getattr(llm_response, "content", None)
149
196
  text = _content_to_text(content)
150
197
  # Keep updating — the last non-empty model reply is the final answer
151
198
  if text:
152
199
  state["output"] = text
153
200
 
201
+ # Token counts for this call
202
+ usage = getattr(llm_response, "usage_metadata", None)
203
+ call_input_tokens = getattr(usage, "prompt_token_count", None) if usage is not None else None
204
+ call_output_tokens = getattr(usage, "candidates_token_count", None) if usage is not None else None
205
+ if call_input_tokens is not None:
206
+ state["input_tokens"] += int(call_input_tokens)
207
+ if call_output_tokens is not None:
208
+ state["output_tokens"] += int(call_output_tokens)
209
+
210
+ # Execution step
211
+ if start_t is not None:
212
+ steps = state["execution_steps"]
213
+ steps.append({
214
+ "name": f"LLM Call {len(steps) + 1}",
215
+ "duration_ms": (end_t - start_t) * 1000,
216
+ "start_time": start_t,
217
+ "end_time": end_t,
218
+ "model": call_start.get("model") if call_start else None,
219
+ "input": call_start.get("input") if call_start else None,
220
+ "output": text,
221
+ "inputTokenSize": call_input_tokens,
222
+ "outputTokenSize": call_output_tokens,
223
+ })
224
+
154
225
  # ------------------------------------------------------------------
155
226
  # Tool callbacks
156
227
  # ------------------------------------------------------------------
@@ -172,15 +243,28 @@ class AgentXADKPlugin(BasePlugin):
172
243
  state = self._runs.get(inv_id)
173
244
  if state is None:
174
245
  return
175
- start = self._tool_starts.pop(id(tool_context), None)
246
+ start_t = self._tool_starts.pop(id(tool_context), None)
247
+ end_t = time.time()
248
+ tool_name = getattr(tool, "name", "unknown")
249
+ tool_input = _safe_serialize(tool_args)
250
+ tool_output = str(result) if result is not None else None
176
251
  tool_call: Dict[str, Any] = {
177
- "name": getattr(tool, "name", "unknown"),
178
- "input": _safe_serialize(tool_args),
179
- "output": str(result)[:500] if result is not None else None,
252
+ "name": tool_name,
253
+ "input": tool_input,
254
+ "output": tool_output,
180
255
  }
181
- if start is not None:
182
- tool_call["latency_ms"] = max(0, int((time.time() - start) * 1000))
256
+ if start_t is not None:
257
+ tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
183
258
  state["tool_calls"].append(tool_call)
259
+ if start_t is not None:
260
+ state["perf_tool_calls"].append({
261
+ "name": tool_name,
262
+ "duration_ms": (end_t - start_t) * 1000,
263
+ "start_time": start_t,
264
+ "end_time": end_t,
265
+ "input": tool_input,
266
+ "output": tool_output,
267
+ })
184
268
 
185
269
  async def on_tool_error_callback(
186
270
  self,
@@ -194,12 +278,25 @@ class AgentXADKPlugin(BasePlugin):
194
278
  state = self._runs.get(inv_id)
195
279
  if state is None:
196
280
  return
197
- start = self._tool_starts.pop(id(tool_context), None)
281
+ start_t = self._tool_starts.pop(id(tool_context), None)
282
+ end_t = time.time()
283
+ tool_name = getattr(tool, "name", "unknown")
284
+ tool_input = _safe_serialize(tool_args)
285
+ tool_output = f"ERROR: {error}"
198
286
  tool_call: Dict[str, Any] = {
199
- "name": getattr(tool, "name", "unknown"),
200
- "input": _safe_serialize(tool_args),
201
- "output": f"ERROR: {error}",
287
+ "name": tool_name,
288
+ "input": tool_input,
289
+ "output": tool_output,
202
290
  }
203
- if start is not None:
204
- tool_call["latency_ms"] = max(0, int((time.time() - start) * 1000))
291
+ if start_t is not None:
292
+ tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
205
293
  state["tool_calls"].append(tool_call)
294
+ if start_t is not None:
295
+ state["perf_tool_calls"].append({
296
+ "name": tool_name,
297
+ "duration_ms": (end_t - start_t) * 1000,
298
+ "start_time": start_t,
299
+ "end_time": end_t,
300
+ "input": tool_input,
301
+ "output": tool_output,
302
+ })