agentx-python 0.6.1__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.6.1/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.6.1 → agentx_python-0.6.2}/agentx/integrations/anthropic.py +54 -25
  4. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/integrations/crewai.py +31 -2
  5. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/integrations/google_adk.py +50 -18
  6. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/integrations/google_genai.py +16 -3
  7. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/integrations/langchain.py +168 -56
  8. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/integrations/openai_agents.py +59 -25
  9. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/tracing/tracer.py +335 -4
  10. agentx_python-0.6.2/agentx/version.py +1 -0
  11. {agentx_python-0.6.1 → agentx_python-0.6.2/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.2}/LICENSE +0 -0
  15. {agentx_python-0.6.1 → agentx_python-0.6.2}/README.md +0 -0
  16. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/__init__.py +0 -0
  17. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/agentx.py +0 -0
  18. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/__init__.py +0 -0
  19. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/_term.py +0 -0
  20. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/adapters/__init__.py +0 -0
  21. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/adapters/http_endpoint.py +0 -0
  22. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/adapters/precomputed.py +0 -0
  23. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/adapters/raw.py +0 -0
  24. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/client.py +0 -0
  25. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/datasets.py +0 -0
  26. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/models.py +0 -0
  27. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/redaction.py +0 -0
  28. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/reporting.py +0 -0
  29. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/results.py +0 -0
  30. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/runner.py +0 -0
  31. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/evaluations/tracing.py +0 -0
  32. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/exceptions.py +0 -0
  33. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/integrations/__init__.py +0 -0
  34. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/resources/__init__.py +0 -0
  35. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/resources/agent.py +0 -0
  36. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/resources/conversation.py +0 -0
  37. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/resources/workforce.py +0 -0
  38. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/tracing/__init__.py +0 -0
  39. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/tracing/ci_types.py +0 -0
  40. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/tracing/ingest_client.py +0 -0
  41. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx/util.py +0 -0
  42. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx_python.egg-info/SOURCES.txt +0 -0
  43. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx_python.egg-info/dependency_links.txt +0 -0
  44. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx_python.egg-info/requires.txt +0 -0
  45. {agentx_python-0.6.1 → agentx_python-0.6.2}/agentx_python.egg-info/top_level.txt +0 -0
  46. {agentx_python-0.6.1 → agentx_python-0.6.2}/setup.cfg +0 -0
  47. {agentx_python-0.6.1 → agentx_python-0.6.2}/setup.py +0 -0
  48. {agentx_python-0.6.1 → 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.6.1
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
+ )
@@ -80,7 +80,7 @@ def _patch_create(
80
80
  try:
81
81
  output = response.content[0].text if response.content else None
82
82
  except Exception:
83
- output = str(response)[:500]
83
+ output = str(response)
84
84
  try:
85
85
  usage = getattr(response, "usage", None)
86
86
  if usage is not None:
@@ -88,30 +88,54 @@ def _patch_create(
88
88
  output_tokens = getattr(usage, "output_tokens", None)
89
89
  except Exception:
90
90
  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
- )
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
+ )
115
139
 
116
140
  patched_create._agentx_patched = True
117
141
  messages_resource.create = patched_create
@@ -162,6 +186,11 @@ def _patch_stream(
162
186
  "duration_ms": latency_ms,
163
187
  "start_time": start_t,
164
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,
165
194
  }],
166
195
  has_errors=error is not None,
167
196
  )
@@ -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(
@@ -52,6 +52,16 @@ def _content_to_text(content: Any) -> Optional[str]:
52
52
  return " ".join(texts) if texts else None
53
53
 
54
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
+
55
65
  class AgentXADKPlugin(BasePlugin):
56
66
  """
57
67
  Google ADK plugin that sends one AgentX trace per runner invocation.
@@ -152,25 +162,31 @@ class AgentXADKPlugin(BasePlugin):
152
162
  ) -> None:
153
163
  inv_id = callback_context.get_invocation_context().invocation_id
154
164
  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.
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.
160
171
  # ADK creates different CallbackContext objects for before vs after, so
161
172
  # id(callback_context) cannot be used as a key across the two calls.
162
173
  if inv_id not in self._model_starts:
163
174
  self._model_starts[inv_id] = []
164
- self._model_starts[inv_id].append(time.time())
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
+ })
165
180
 
166
181
  async def after_model_callback(
167
182
  self, *, callback_context: Any, llm_response: Any
168
183
  ) -> None:
169
184
  inv_id = callback_context.get_invocation_context().invocation_id
170
185
  state = self._runs.get(inv_id)
171
- # Pop the earliest queued start time (FIFO — model calls are sequential)
186
+ # Pop the earliest queued call (FIFO — model calls are sequential)
172
187
  starts = self._model_starts.get(inv_id, [])
173
- start_t = starts.pop(0) if starts else None
188
+ call_start = starts.pop(0) if starts else None
189
+ start_t = call_start.get("start") if call_start else None
174
190
  end_t = time.time()
175
191
 
176
192
  if state is None:
@@ -182,6 +198,15 @@ class AgentXADKPlugin(BasePlugin):
182
198
  if text:
183
199
  state["output"] = text
184
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
+
185
210
  # Execution step
186
211
  if start_t is not None:
187
212
  steps = state["execution_steps"]
@@ -190,14 +215,13 @@ class AgentXADKPlugin(BasePlugin):
190
215
  "duration_ms": (end_t - start_t) * 1000,
191
216
  "start_time": start_t,
192
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,
193
223
  })
194
224
 
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
225
  # ------------------------------------------------------------------
202
226
  # Tool callbacks
203
227
  # ------------------------------------------------------------------
@@ -222,10 +246,12 @@ class AgentXADKPlugin(BasePlugin):
222
246
  start_t = self._tool_starts.pop(id(tool_context), None)
223
247
  end_t = time.time()
224
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
225
251
  tool_call: Dict[str, Any] = {
226
252
  "name": tool_name,
227
- "input": _safe_serialize(tool_args),
228
- "output": str(result)[:500] if result is not None else None,
253
+ "input": tool_input,
254
+ "output": tool_output,
229
255
  }
230
256
  if start_t is not None:
231
257
  tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
@@ -236,6 +262,8 @@ class AgentXADKPlugin(BasePlugin):
236
262
  "duration_ms": (end_t - start_t) * 1000,
237
263
  "start_time": start_t,
238
264
  "end_time": end_t,
265
+ "input": tool_input,
266
+ "output": tool_output,
239
267
  })
240
268
 
241
269
  async def on_tool_error_callback(
@@ -253,10 +281,12 @@ class AgentXADKPlugin(BasePlugin):
253
281
  start_t = self._tool_starts.pop(id(tool_context), None)
254
282
  end_t = time.time()
255
283
  tool_name = getattr(tool, "name", "unknown")
284
+ tool_input = _safe_serialize(tool_args)
285
+ tool_output = f"ERROR: {error}"
256
286
  tool_call: Dict[str, Any] = {
257
287
  "name": tool_name,
258
- "input": _safe_serialize(tool_args),
259
- "output": f"ERROR: {error}",
288
+ "input": tool_input,
289
+ "output": tool_output,
260
290
  }
261
291
  if start_t is not None:
262
292
  tool_call["latency_ms"] = max(0, int((end_t - start_t) * 1000))
@@ -267,4 +297,6 @@ class AgentXADKPlugin(BasePlugin):
267
297
  "duration_ms": (end_t - start_t) * 1000,
268
298
  "start_time": start_t,
269
299
  "end_time": end_t,
300
+ "input": tool_input,
301
+ "output": tool_output,
270
302
  })
@@ -100,6 +100,7 @@ def _patch_generate_content(
100
100
  if usage is not None:
101
101
  input_tokens = getattr(usage, "prompt_token_count", None)
102
102
  output_tokens = getattr(usage, "candidates_token_count", None)
103
+ input_repr = contents if isinstance(contents, str) else _safe_serialize(contents)
103
104
  perf = build_performance_summary(
104
105
  total_duration_ms=latency_ms,
105
106
  execution_steps=[{
@@ -107,12 +108,17 @@ def _patch_generate_content(
107
108
  "duration_ms": latency_ms,
108
109
  "start_time": start_t,
109
110
  "end_time": end_t,
111
+ "model": str(model) if model else None,
112
+ "input": input_repr,
113
+ "output": output,
114
+ "inputTokenSize": input_tokens,
115
+ "outputTokenSize": output_tokens,
110
116
  }],
111
117
  has_errors=error is not None,
112
118
  )
113
119
  tracer._send(
114
120
  name=name,
115
- input=contents if isinstance(contents, str) else _safe_serialize(contents),
121
+ input=input_repr,
116
122
  output=output,
117
123
  latency_ms=latency_ms,
118
124
  error=error,
@@ -169,6 +175,8 @@ def _patch_generate_content_stream(
169
175
  if last_usage_metadata is not None:
170
176
  input_tokens = getattr(last_usage_metadata, "prompt_token_count", None)
171
177
  output_tokens = getattr(last_usage_metadata, "candidates_token_count", None)
178
+ input_repr = contents if isinstance(contents, str) else _safe_serialize(contents)
179
+ output_repr = "".join(accumulated_text) or None
172
180
  perf = build_performance_summary(
173
181
  total_duration_ms=latency_ms,
174
182
  execution_steps=[{
@@ -176,13 +184,18 @@ def _patch_generate_content_stream(
176
184
  "duration_ms": latency_ms,
177
185
  "start_time": start_t,
178
186
  "end_time": end_t,
187
+ "model": str(model) if model else None,
188
+ "input": input_repr,
189
+ "output": output_repr,
190
+ "inputTokenSize": input_tokens,
191
+ "outputTokenSize": output_tokens,
179
192
  }],
180
193
  has_errors=error is not None,
181
194
  )
182
195
  tracer._send(
183
196
  name=name,
184
- input=contents if isinstance(contents, str) else _safe_serialize(contents),
185
- output="".join(accumulated_text) or None,
197
+ input=input_repr,
198
+ output=output_repr,
186
199
  latency_ms=latency_ms,
187
200
  error=error,
188
201
  framework="google-genai",