2828
2929from . import _metrics
3030from . import tracing
31- from ..events import event as event_lib
3231from ._schema_version import resolve_schema_version
3332from ._schema_version import SCHEMA_VERSION_SEMCONV_ALIGNED
3433
34+ # pylint: disable=g-import-not-at-top
3535if TYPE_CHECKING :
3636 from ..agents .base_agent import BaseAgent
3737 from ..agents .invocation_context import InvocationContext
38+ from ..events import event as event_lib
3839 from ..models .llm_request import LlmRequest
3940 from ..models .llm_response import LlmResponse
4041 from ..tools .base_tool import BaseTool
4142 from ..workflow ._base_node import BaseNode
4243
4344logger = logging .getLogger ("google_adk." + __name__ )
4445
46+ _INVOKE_AGENT_TELEMETRY_KEY = context_api .create_key ("invoke_agent_telemetry" )
47+
4548
4649@contextlib .contextmanager
4750def record_invocation (
@@ -91,6 +94,22 @@ class TelemetryContext:
9194 error_type : str | None = None
9295 span : tracing .GenerateContentSpan | trace .Span | None = None
9396 _llm_responses : list [LlmResponse ] = dataclasses .field (default_factory = list )
97+ _inference_call_count : int = 0
98+ _tool_call_count : int = 0
99+
100+ @property
101+ def inference_call_count (self ) -> int :
102+ return self ._inference_call_count
103+
104+ def increment_inference_calls (self ) -> None :
105+ self ._inference_call_count += 1
106+
107+ @property
108+ def tool_call_count (self ) -> int :
109+ return self ._tool_call_count
110+
111+ def increment_tool_calls (self ) -> None :
112+ self ._tool_call_count += 1
94113
95114 @property
96115 def llm_responses (self ) -> list [LlmResponse ]:
@@ -120,6 +139,36 @@ def _record_agent_metrics(
120139 logger .exception ("Failed to record agent metrics for agent %s" , agent_name )
121140
122141
142+ def _flush_invoke_agent_metrics (
143+ tel_ctx : TelemetryContext , agent_name : str
144+ ) -> None :
145+ """Flushes this span's accumulated inference/tool-call metrics."""
146+ _metrics .record_invoke_agent_inference_calls (
147+ agent_name , tel_ctx .inference_call_count
148+ )
149+ _metrics .record_invoke_agent_tool_calls (agent_name , tel_ctx .tool_call_count )
150+
151+
152+ def _active_invoke_agent_tel_ctx () -> TelemetryContext | None :
153+ """Returns the TelemetryContext of the active invoke_agent span."""
154+ value = context_api .get_value (_INVOKE_AGENT_TELEMETRY_KEY )
155+ return value if isinstance (value , TelemetryContext ) else None
156+
157+
158+ def _accumulate_invoke_agent_tool_call () -> None :
159+ """Counts one tool call against the active invoke_agent span."""
160+ span_tel_ctx = _active_invoke_agent_tel_ctx ()
161+ if span_tel_ctx is not None :
162+ span_tel_ctx .increment_tool_calls ()
163+
164+
165+ def _accumulate_invoke_agent_inference_call () -> None :
166+ """Counts one model call against the active invoke_agent span."""
167+ span_tel_ctx = _active_invoke_agent_tel_ctx ()
168+ if span_tel_ctx is not None :
169+ span_tel_ctx .increment_inference_calls ()
170+
171+
123172@contextlib .asynccontextmanager
124173async def record_agent_invocation (
125174 ctx : InvocationContext , agent : BaseAgent
@@ -129,30 +178,36 @@ async def record_agent_invocation(
129178 caught_error : Exception | None = None
130179 span : trace .Span | None = None
131180 span_name = f"invoke_agent { agent .name } "
181+ tel_ctx = TelemetryContext ()
182+ token = context_api .attach (
183+ context_api .set_value (_INVOKE_AGENT_TELEMETRY_KEY , tel_ctx )
184+ )
132185 try :
133186 with tracing .tracer .start_as_current_span (span_name ) as s :
134187 span = s
135188 tracing .trace_agent_invocation (span , agent , ctx )
136- tel_ctx = TelemetryContext ( otel_context = context_api .get_current () )
189+ tel_ctx . otel_context = context_api .get_current ()
137190 yield tel_ctx
138191 except Exception as e :
139192 caught_error = e
140193 raise
141194 finally :
195+ context_api .detach (token )
142196 _record_agent_metrics (
143197 agent .name ,
144198 _metrics .get_elapsed_s (span , start_time ),
145199 getattr (getattr (ctx , "session" , None ), "events" , []),
146200 caught_error ,
147201 )
202+ _flush_invoke_agent_metrics (tel_ctx , agent .name )
148203
149204
150205@contextlib .asynccontextmanager
151206async def record_tool_execution (
152207 tool : BaseTool ,
153208 agent : BaseAgent ,
154209 function_args : dict [str , object ],
155- invocation_context : InvocationContext | None = None ,
210+ invocation_context : InvocationContext ,
156211) -> AsyncIterator [TelemetryContext ]:
157212 """Unified context manager for consolidated tool execution telemetry."""
158213 start_time = time .monotonic ()
@@ -181,6 +236,7 @@ async def record_tool_execution(
181236 error_type = tel_ctx .error_type ,
182237 )
183238 finally :
239+ _accumulate_invoke_agent_tool_call ()
184240 try :
185241 _metrics .record_tool_execution_duration (
186242 tool_name = tool .name ,
@@ -214,6 +270,7 @@ async def record_inference_telemetry(
214270 yield tel_ctx
215271 finally :
216272 inference_error = sys .exc_info ()[1 ]
273+ _accumulate_invoke_agent_inference_call ()
217274 agent = invocation_context .agent
218275 elapsed_s = _metrics .get_elapsed_s (tel_ctx .span , start_time )
219276 try :
0 commit comments