2727
2828from . import _metrics
2929from . import tracing
30+ from ..events import event as event_lib
3031
3132if TYPE_CHECKING :
3233 from ..agents .base_agent import BaseAgent
3334 from ..agents .invocation_context import InvocationContext
34- from ..events import event as event_lib
3535 from ..models .llm_request import LlmRequest
3636 from ..models .llm_response import LlmResponse
3737 from ..tools .base_tool import BaseTool
@@ -78,31 +78,11 @@ class TelemetryContext:
7878 error_type : str | None = None
7979 span : tracing .GenerateContentSpan | trace .Span | None = None
8080 _llm_responses : list [LlmResponse ] = dataclasses .field (default_factory = list )
81- _inference_call_count : int = 0
82- _tool_call_count : int = 0
8381
8482 @property
8583 def llm_responses (self ) -> list [LlmResponse ]:
8684 return self ._llm_responses
8785
88- @property
89- def inference_call_count (self ) -> int :
90- """Number of model calls counted against this invoke_agent span."""
91- return self ._inference_call_count
92-
93- def increment_inference_calls (self ) -> None :
94- """Counts one model call against this span (including calls that raised)."""
95- self ._inference_call_count += 1
96-
97- def increment_tool_calls (self ) -> None :
98- """Counts one tool call against this invoke_agent span."""
99- self ._tool_call_count += 1
100-
101- @property
102- def tool_call_count (self ) -> int :
103- """Number of tool calls counted against this invoke_agent span."""
104- return self ._tool_call_count
105-
10686 def record_llm_response (
10787 self , invocation_context : InvocationContext , response : LlmResponse
10888 ) -> None :
@@ -130,30 +110,6 @@ def _record_agent_metrics(
130110 logger .exception ("Failed to record agent metrics for agent %s" , agent_name )
131111
132112
133- def _flush_invoke_agent_metrics (
134- tel_ctx : TelemetryContext , agent_name : str
135- ) -> None :
136- """Flushes this span's accumulated inference/tool-call metrics."""
137- _metrics .record_invoke_agent_inference_calls (
138- agent_name , tel_ctx .inference_call_count
139- )
140- _metrics .record_invoke_agent_tool_calls (agent_name , tel_ctx .tool_call_count )
141-
142-
143- def _accumulate_invoke_agent_tool_call (ctx : InvocationContext ) -> None :
144- """Counts one tool call against the active invoke_agent span."""
145- span_tel_ctx = ctx ._invoke_agent_telemetry_context # pylint: disable=protected-access
146- if span_tel_ctx is not None :
147- span_tel_ctx .increment_tool_calls ()
148-
149-
150- def _accumulate_invoke_agent_inference_call (ctx : InvocationContext ) -> None :
151- """Counts one model call against the active invoke_agent span."""
152- span_tel_ctx = ctx ._invoke_agent_telemetry_context # pylint: disable=protected-access
153- if span_tel_ctx is not None :
154- span_tel_ctx .increment_inference_calls ()
155-
156-
157113@contextlib .asynccontextmanager
158114async def record_agent_invocation (
159115 ctx : InvocationContext , agent : BaseAgent
@@ -163,13 +119,11 @@ async def record_agent_invocation(
163119 caught_error : Exception | None = None
164120 span : trace .Span | None = None
165121 span_name = f"invoke_agent { agent .name } "
166- tel_ctx = TelemetryContext ()
167122 try :
168123 with tracing .tracer .start_as_current_span (span_name ) as s :
169124 span = s
170125 tracing .trace_agent_invocation (span , agent , ctx )
171- tel_ctx .otel_context = context_api .get_current ()
172- ctx ._invoke_agent_telemetry_context = tel_ctx # pylint: disable=protected-access
126+ tel_ctx = TelemetryContext (otel_context = context_api .get_current ())
173127 yield tel_ctx
174128 except Exception as e :
175129 caught_error = e
@@ -182,15 +136,14 @@ async def record_agent_invocation(
182136 getattr (getattr (ctx , "session" , None ), "events" , []),
183137 caught_error ,
184138 )
185- _flush_invoke_agent_metrics (tel_ctx , agent .name )
186139
187140
188141@contextlib .asynccontextmanager
189142async def record_tool_execution (
190143 tool : BaseTool ,
191144 agent : BaseAgent ,
192145 function_args : dict [str , object ],
193- invocation_context : InvocationContext ,
146+ invocation_context : InvocationContext | None = None ,
194147) -> AsyncIterator [TelemetryContext ]:
195148 """Unified context manager for consolidated tool execution telemetry."""
196149 start_time = time .monotonic ()
@@ -218,7 +171,6 @@ async def record_tool_execution(
218171 invocation_context = invocation_context ,
219172 error_type = tel_ctx .error_type ,
220173 )
221- _accumulate_invoke_agent_tool_call (invocation_context )
222174 finally :
223175 try :
224176 _metrics .record_tool_execution_duration (
@@ -253,7 +205,6 @@ async def record_inference_telemetry(
253205 yield tel_ctx
254206 finally :
255207 inference_error = sys .exc_info ()[1 ]
256- _accumulate_invoke_agent_inference_call (invocation_context )
257208 agent = invocation_context .agent
258209 elapsed_s = _get_elapsed_s (tel_ctx .span , start_time )
259210 try :
0 commit comments