agentx-python 0.6.1__tar.gz → 0.6.3__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.6.1/agentx_python.egg-info → agentx_python-0.6.3}/PKG-INFO +1 -1
  2. agentx_python-0.6.3/agentx/integrations/_perf.py +226 -0
  3. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/anthropic.py +83 -29
  4. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/crewai.py +31 -2
  5. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/google_adk.py +67 -20
  6. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/google_genai.py +30 -5
  7. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/langchain.py +200 -56
  8. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/openai_agents.py +59 -25
  9. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/tracing/tracer.py +354 -4
  10. agentx_python-0.6.3/agentx/version.py +1 -0
  11. {agentx_python-0.6.1 → agentx_python-0.6.3/agentx_python.egg-info}/PKG-INFO +1 -1
  12. agentx_python-0.6.1/agentx/integrations/_perf.py +0 -116
  13. agentx_python-0.6.1/agentx/version.py +0 -1
  14. {agentx_python-0.6.1 → agentx_python-0.6.3}/LICENSE +0 -0
  15. {agentx_python-0.6.1 → agentx_python-0.6.3}/README.md +0 -0
  16. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/__init__.py +0 -0
  17. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/agentx.py +0 -0
  18. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/__init__.py +0 -0
  19. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/_term.py +0 -0
  20. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/adapters/__init__.py +0 -0
  21. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/adapters/http_endpoint.py +0 -0
  22. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/adapters/precomputed.py +0 -0
  23. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/adapters/raw.py +0 -0
  24. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/client.py +0 -0
  25. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/datasets.py +0 -0
  26. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/models.py +0 -0
  27. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/redaction.py +0 -0
  28. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/reporting.py +0 -0
  29. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/results.py +0 -0
  30. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/runner.py +0 -0
  31. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/evaluations/tracing.py +0 -0
  32. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/exceptions.py +0 -0
  33. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/integrations/__init__.py +0 -0
  34. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/resources/__init__.py +0 -0
  35. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/resources/agent.py +0 -0
  36. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/resources/conversation.py +0 -0
  37. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/resources/workforce.py +0 -0
  38. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/tracing/__init__.py +0 -0
  39. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/tracing/ci_types.py +0 -0
  40. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/tracing/ingest_client.py +0 -0
  41. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx/util.py +0 -0
  42. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx_python.egg-info/SOURCES.txt +0 -0
  43. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx_python.egg-info/dependency_links.txt +0 -0
  44. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx_python.egg-info/requires.txt +0 -0
  45. {agentx_python-0.6.1 → agentx_python-0.6.3}/agentx_python.egg-info/top_level.txt +0 -0
  46. {agentx_python-0.6.1 → agentx_python-0.6.3}/setup.cfg +0 -0
  47. {agentx_python-0.6.1 → agentx_python-0.6.3}/setup.py +0 -0
  48. {agentx_python-0.6.1 → agentx_python-0.6.3}/tests/test_integration.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentx-python
3
- Version: 0.6.1
3
+ Version: 0.6.3
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
+ )
@@ -22,6 +22,34 @@ from agentx.tracing.tracer import Tracer, _safe_serialize
22
22
  from agentx.integrations._perf import build_performance_summary
23
23
 
24
24
 
25
+ def _extract_output_text(response: Any) -> Optional[str]:
26
+ """
27
+ Extract the assistant's text reply from a Messages API response, falling
28
+ back to a description of any tool_use blocks when the response is a pure
29
+ tool call with no accompanying text.
30
+ """
31
+ content = getattr(response, "content", None) if response is not None else None
32
+ if not content:
33
+ return None
34
+ texts = []
35
+ tool_calls = []
36
+ for block in content:
37
+ block_type = getattr(block, "type", None)
38
+ if block_type == "text":
39
+ text = getattr(block, "text", None)
40
+ if text:
41
+ texts.append(text)
42
+ elif block_type == "tool_use":
43
+ name = getattr(block, "name", "unknown")
44
+ tool_input = getattr(block, "input", None)
45
+ tool_calls.append(f"{name}({tool_input})")
46
+ if texts:
47
+ return "\n".join(texts)
48
+ if tool_calls:
49
+ return "[tool call] " + ", ".join(tool_calls)
50
+ return None
51
+
52
+
25
53
  def patch_anthropic_client(
26
54
  client: Any,
27
55
  tracer: Tracer,
@@ -77,10 +105,7 @@ def _patch_create(
77
105
  input_tokens = None
78
106
  output_tokens = None
79
107
  if response is not None:
80
- try:
81
- output = response.content[0].text if response.content else None
82
- except Exception:
83
- output = str(response)[:500]
108
+ output = _extract_output_text(response)
84
109
  try:
85
110
  usage = getattr(response, "usage", None)
86
111
  if usage is not None:
@@ -88,30 +113,54 @@ def _patch_create(
88
113
  output_tokens = getattr(usage, "output_tokens", None)
89
114
  except Exception:
90
115
  pass
91
- perf = build_performance_summary(
92
- total_duration_ms=latency_ms,
93
- execution_steps=[{
94
- "name": "LLM Call 1",
95
- "duration_ms": latency_ms,
96
- "start_time": start_t,
97
- "end_time": end_t,
98
- }],
99
- has_errors=error is not None,
100
- )
101
- tracer._send(
102
- name=name,
103
- input=_safe_serialize(input_messages),
104
- output=output,
105
- latency_ms=latency_ms,
106
- error=error,
107
- framework="anthropic",
108
- model=model,
109
- metadata=metadata,
110
- session_id=session_id,
111
- performance_summary=perf,
112
- input_tokens=input_tokens,
113
- output_tokens=output_tokens,
114
- )
116
+
117
+ active_span = tracer.current_span
118
+ if active_span is not None:
119
+ # Part of a `with tracer.trace(...)` block (e.g. a multi-call
120
+ # agentic loop) — attach as one LLM-call step on that span's
121
+ # trace instead of sending an independent trace per call.
122
+ if error is not None:
123
+ active_span.set_error(error)
124
+ active_span._record_llm_call(
125
+ duration_ms=latency_ms,
126
+ start_time=start_t,
127
+ end_time=end_t,
128
+ input=_safe_serialize(input_messages),
129
+ output=output,
130
+ model=model,
131
+ input_tokens=input_tokens,
132
+ output_tokens=output_tokens,
133
+ )
134
+ else:
135
+ perf = build_performance_summary(
136
+ total_duration_ms=latency_ms,
137
+ execution_steps=[{
138
+ "name": "LLM Call 1",
139
+ "duration_ms": latency_ms,
140
+ "start_time": start_t,
141
+ "end_time": end_t,
142
+ "model": model,
143
+ "input": _safe_serialize(input_messages),
144
+ "output": output,
145
+ "inputTokenSize": input_tokens,
146
+ "outputTokenSize": output_tokens,
147
+ }],
148
+ has_errors=error is not None,
149
+ )
150
+ tracer._send(
151
+ name=name,
152
+ input=_safe_serialize(input_messages),
153
+ output=output,
154
+ latency_ms=latency_ms,
155
+ error=error,
156
+ framework="anthropic",
157
+ model=model,
158
+ metadata=metadata,
159
+ session_id=session_id,
160
+ performance_summary=perf,
161
+ input_tokens=input_tokens,
162
+ output_tokens=output_tokens,
163
+ )
115
164
 
116
165
  patched_create._agentx_patched = True
117
166
  messages_resource.create = patched_create
@@ -148,7 +197,7 @@ def _patch_stream(
148
197
  output_tokens = None
149
198
  try:
150
199
  final = ctx.get_final_message()
151
- output = final.content[0].text if final.content else None
200
+ output = _extract_output_text(final)
152
201
  usage = getattr(final, "usage", None)
153
202
  if usage is not None:
154
203
  input_tokens = getattr(usage, "input_tokens", None)
@@ -162,6 +211,11 @@ def _patch_stream(
162
211
  "duration_ms": latency_ms,
163
212
  "start_time": start_t,
164
213
  "end_time": end_t,
214
+ "model": kwargs.get("model"),
215
+ "input": _safe_serialize(kwargs.get("messages")),
216
+ "output": output,
217
+ "inputTokenSize": input_tokens,
218
+ "outputTokenSize": output_tokens,
165
219
  }],
166
220
  has_errors=error is not None,
167
221
  )
@@ -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(
@@ -36,7 +36,11 @@ except ImportError as exc: # pragma: no cover
36
36
 
37
37
 
38
38
  def _content_to_text(content: Any) -> Optional[str]:
39
- """Extract plain text from a google.genai types.Content object."""
39
+ """
40
+ Extract plain text from a google.genai types.Content object, falling back
41
+ to a description of any function_call parts when there's no text (Gemini
42
+ function calling — the model responded with a pure tool call).
43
+ """
40
44
  if content is None:
41
45
  return None
42
46
  if isinstance(content, str):
@@ -45,11 +49,32 @@ def _content_to_text(content: Any) -> Optional[str]:
45
49
  if not parts:
46
50
  return None
47
51
  texts = []
52
+ function_calls = []
48
53
  for part in parts:
49
54
  text = getattr(part, "text", None)
50
55
  if text and isinstance(text, str):
51
56
  texts.append(text)
52
- return " ".join(texts) if texts else None
57
+ continue
58
+ fc = getattr(part, "function_call", None)
59
+ if fc is not None:
60
+ name = getattr(fc, "name", "unknown")
61
+ args = getattr(fc, "args", None)
62
+ function_calls.append(f"{name}({args})")
63
+ if texts:
64
+ return " ".join(texts)
65
+ if function_calls:
66
+ return "[tool call] " + ", ".join(function_calls)
67
+ return None
68
+
69
+
70
+ def _contents_to_text(contents: Any) -> Optional[str]:
71
+ """Extract plain text from an LlmRequest's ``contents`` (a list of Content)."""
72
+ if contents is None:
73
+ return None
74
+ if not isinstance(contents, (list, tuple)):
75
+ contents = [contents]
76
+ texts = [t for t in (_content_to_text(c) for c in contents) if t]
77
+ return "\n".join(texts) if texts else None
53
78
 
54
79
 
55
80
  class AgentXADKPlugin(BasePlugin):
@@ -152,25 +177,31 @@ class AgentXADKPlugin(BasePlugin):
152
177
  ) -> None:
153
178
  inv_id = callback_context.get_invocation_context().invocation_id
154
179
  state = self._runs.get(inv_id)
155
- if state and not state["model"]:
156
- model = getattr(llm_request, "model", None)
157
- if model:
158
- state["model"] = str(model)
159
- # Push start time onto the per-invocation stack.
180
+ model = getattr(llm_request, "model", None)
181
+ model_str = str(model) if model else None
182
+ if state and not state["model"] and model_str:
183
+ state["model"] = model_str
184
+ # Push start time + this call's model/input onto the per-invocation
185
+ # stack, so after_model_callback can pair them back up.
160
186
  # ADK creates different CallbackContext objects for before vs after, so
161
187
  # id(callback_context) cannot be used as a key across the two calls.
162
188
  if inv_id not in self._model_starts:
163
189
  self._model_starts[inv_id] = []
164
- self._model_starts[inv_id].append(time.time())
190
+ self._model_starts[inv_id].append({
191
+ "start": time.time(),
192
+ "model": model_str,
193
+ "input": _contents_to_text(getattr(llm_request, "contents", None)),
194
+ })
165
195
 
166
196
  async def after_model_callback(
167
197
  self, *, callback_context: Any, llm_response: Any
168
198
  ) -> None:
169
199
  inv_id = callback_context.get_invocation_context().invocation_id
170
200
  state = self._runs.get(inv_id)
171
- # Pop the earliest queued start time (FIFO — model calls are sequential)
201
+ # Pop the earliest queued call (FIFO — model calls are sequential)
172
202
  starts = self._model_starts.get(inv_id, [])
173
- start_t = starts.pop(0) if starts else None
203
+ call_start = starts.pop(0) if starts else None
204
+ start_t = call_start.get("start") if call_start else None
174
205
  end_t = time.time()
175
206
 
176
207
  if state is None:
@@ -182,6 +213,15 @@ class AgentXADKPlugin(BasePlugin):
182
213
  if text:
183
214
  state["output"] = text
184
215
 
216
+ # Token counts for this call
217
+ usage = getattr(llm_response, "usage_metadata", None)
218
+ call_input_tokens = getattr(usage, "prompt_token_count", None) if usage is not None else None
219
+ call_output_tokens = getattr(usage, "candidates_token_count", None) if usage is not None else None
220
+ if call_input_tokens is not None:
221
+ state["input_tokens"] += int(call_input_tokens)
222
+ if call_output_tokens is not None:
223
+ state["output_tokens"] += int(call_output_tokens)
224
+
185
225
  # Execution step
186
226
  if start_t is not None:
187
227
  steps = state["execution_steps"]
@@ -190,14 +230,13 @@ class AgentXADKPlugin(BasePlugin):
190
230
  "duration_ms": (end_t - start_t) * 1000,
191
231
  "start_time": start_t,
192
232
  "end_time": end_t,
233
+ "model": call_start.get("model") if call_start else None,
234
+ "input": call_start.get("input") if call_start else None,
235
+ "output": text,
236
+ "inputTokenSize": call_input_tokens,
237
+ "outputTokenSize": call_output_tokens,
193
238
  })
194
239
 
195
- # Token counts
196
- usage = getattr(llm_response, "usage_metadata", None)
197
- if usage is not None:
198
- state["input_tokens"] += int(getattr(usage, "prompt_token_count", None) or 0)
199
- state["output_tokens"] += int(getattr(usage, "candidates_token_count", None) or 0)
200
-
201
240
  # ------------------------------------------------------------------
202
241
  # Tool callbacks
203
242
  # ------------------------------------------------------------------
@@ -222,10 +261,12 @@ class AgentXADKPlugin(BasePlugin):
222
261
  start_t = self._tool_starts.pop(id(tool_context), None)
223
262
  end_t = time.time()
224
263
  tool_name = getattr(tool, "name", "unknown")
264
+ tool_input = _safe_serialize(tool_args)
265
+ tool_output = str(result) if result is not None else None
225
266
  tool_call: Dict[str, Any] = {
226
267
  "name": tool_name,
227
- "input": _safe_serialize(tool_args),
228
- "output": str(result)[:500] if result is not None else None,
268
+ "input": tool_input,
269
+ "output": tool_output,
229
270
  }
230
271
  if start_t is not None:
231
272
  tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
@@ -236,6 +277,8 @@ class AgentXADKPlugin(BasePlugin):
236
277
  "duration_ms": (end_t - start_t) * 1000,
237
278
  "start_time": start_t,
238
279
  "end_time": end_t,
280
+ "input": tool_input,
281
+ "output": tool_output,
239
282
  })
240
283
 
241
284
  async def on_tool_error_callback(
@@ -253,10 +296,12 @@ class AgentXADKPlugin(BasePlugin):
253
296
  start_t = self._tool_starts.pop(id(tool_context), None)
254
297
  end_t = time.time()
255
298
  tool_name = getattr(tool, "name", "unknown")
299
+ tool_input = _safe_serialize(tool_args)
300
+ tool_output = f"ERROR: {error}"
256
301
  tool_call: Dict[str, Any] = {
257
302
  "name": tool_name,
258
- "input": _safe_serialize(tool_args),
259
- "output": f"ERROR: {error}",
303
+ "input": tool_input,
304
+ "output": tool_output,
260
305
  }
261
306
  if start_t is not None:
262
307
  tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
@@ -267,4 +312,6 @@ class AgentXADKPlugin(BasePlugin):
267
312
  "duration_ms": (end_t - start_t) * 1000,
268
313
  "start_time": start_t,
269
314
  "end_time": end_t,
315
+ "input": tool_input,
316
+ "output": tool_output,
270
317
  })