[pull] main from google:main

src/google/adk/telemetry/_instrumentation.py

@@ -28,20 +28,23 @@
 
 from . import _metrics
 from . import tracing
-from ..events import event as event_lib
 from ._schema_version import resolve_schema_version
 from ._schema_version import SCHEMA_VERSION_SEMCONV_ALIGNED
 
+# pylint: disable=g-import-not-at-top
 if TYPE_CHECKING:
   from ..agents.base_agent import BaseAgent
   from ..agents.invocation_context import InvocationContext
+  from ..events import event as event_lib
   from ..models.llm_request import LlmRequest
   from ..models.llm_response import LlmResponse
   from ..tools.base_tool import BaseTool
   from ..workflow._base_node import BaseNode
 
 logger = logging.getLogger("google_adk." + __name__)
 
+_INVOKE_AGENT_TELEMETRY_KEY = context_api.create_key("invoke_agent_telemetry")
+
 
 @contextlib.contextmanager
 def record_invocation(
@@ -91,6 +94,22 @@ class TelemetryContext:
   error_type: str | None = None
   span: tracing.GenerateContentSpan | trace.Span | None = None
   _llm_responses: list[LlmResponse] = dataclasses.field(default_factory=list)
+  _inference_call_count: int = 0
+  _tool_call_count: int = 0
+
+  @property
+  def inference_call_count(self) -> int:
+    return self._inference_call_count
+
+  def increment_inference_calls(self) -> None:
+    self._inference_call_count += 1
+
+  @property
+  def tool_call_count(self) -> int:
+    return self._tool_call_count
+
+  def increment_tool_calls(self) -> None:
+    self._tool_call_count += 1
 
   @property
   def llm_responses(self) -> list[LlmResponse]:
@@ -120,6 +139,36 @@ def _record_agent_metrics(
     logger.exception("Failed to record agent metrics for agent %s", agent_name)
 
 
+def _flush_invoke_agent_metrics(
+    tel_ctx: TelemetryContext, agent_name: str
+) -> None:
+  """Flushes this span's accumulated inference/tool-call metrics."""
+  _metrics.record_invoke_agent_inference_calls(
+      agent_name, tel_ctx.inference_call_count
+  )
+  _metrics.record_invoke_agent_tool_calls(agent_name, tel_ctx.tool_call_count)
+
+
+def _active_invoke_agent_tel_ctx() -> TelemetryContext | None:
+  """Returns the TelemetryContext of the active invoke_agent span."""
+  value = context_api.get_value(_INVOKE_AGENT_TELEMETRY_KEY)
+  return value if isinstance(value, TelemetryContext) else None
+
+
+def _accumulate_invoke_agent_tool_call() -> None:
+  """Counts one tool call against the active invoke_agent span."""
+  span_tel_ctx = _active_invoke_agent_tel_ctx()
+  if span_tel_ctx is not None:
+    span_tel_ctx.increment_tool_calls()
+
+
+def _accumulate_invoke_agent_inference_call() -> None:
+  """Counts one model call against the active invoke_agent span."""
+  span_tel_ctx = _active_invoke_agent_tel_ctx()
+  if span_tel_ctx is not None:
+    span_tel_ctx.increment_inference_calls()
+
+
 @contextlib.asynccontextmanager
 async def record_agent_invocation(
     ctx: InvocationContext, agent: BaseAgent
@@ -129,30 +178,36 @@ async def record_agent_invocation(
   caught_error: Exception | None = None
   span: trace.Span | None = None
   span_name = f"invoke_agent {agent.name}"
+  tel_ctx = TelemetryContext()
+  token = context_api.attach(
+      context_api.set_value(_INVOKE_AGENT_TELEMETRY_KEY, tel_ctx)
+  )
   try:
     with tracing.tracer.start_as_current_span(span_name) as s:
       span = s
       tracing.trace_agent_invocation(span, agent, ctx)
-      tel_ctx = TelemetryContext(otel_context=context_api.get_current())
+      tel_ctx.otel_context = context_api.get_current()
       yield tel_ctx
   except Exception as e:
     caught_error = e
     raise
   finally:
+    context_api.detach(token)
     _record_agent_metrics(
         agent.name,
         _metrics.get_elapsed_s(span, start_time),
         getattr(getattr(ctx, "session", None), "events", []),
         caught_error,
     )
+    _flush_invoke_agent_metrics(tel_ctx, agent.name)
 
 
 @contextlib.asynccontextmanager
 async def record_tool_execution(
     tool: BaseTool,
     agent: BaseAgent,
     function_args: dict[str, object],
-    invocation_context: InvocationContext | None = None,
+    invocation_context: InvocationContext,
 ) -> AsyncIterator[TelemetryContext]:
   """Unified context manager for consolidated tool execution telemetry."""
   start_time = time.monotonic()
@@ -181,6 +236,7 @@ async def record_tool_execution(
             error_type=tel_ctx.error_type,
         )
   finally:
+    _accumulate_invoke_agent_tool_call()
     try:
       _metrics.record_tool_execution_duration(
           tool_name=tool.name,
@@ -214,6 +270,7 @@ async def record_inference_telemetry(
       yield tel_ctx
   finally:
     inference_error = sys.exc_info()[1]
+    _accumulate_invoke_agent_inference_call()
     agent = invocation_context.agent
     elapsed_s = _metrics.get_elapsed_s(tel_ctx.span, start_time)
     try:

src/google/adk/telemetry/_metrics.py

@@ -113,6 +113,44 @@
     gen_ai_metrics.create_gen_ai_client_operation_duration(meter)
 )
 _client_token_usage = gen_ai_metrics.create_gen_ai_client_token_usage(meter)
+_invoke_agent_inference_calls = meter.create_histogram(
+    "gen_ai.invoke_agent.inference_calls",
+    unit="1",
+    description="Number of inference (model) calls per agent invocation.",
+    explicit_bucket_boundaries_advisory=[
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        8,
+        12,
+        16,
+        24,
+        32,
+        64,
+    ],
+)
+_invoke_agent_tool_calls = meter.create_histogram(
+    "gen_ai.invoke_agent.tool_calls",
+    unit="1",
+    description="Number of tool calls per agent invocation.",
+    explicit_bucket_boundaries_advisory=[
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        8,
+        12,
+        16,
+        24,
+        32,
+        64,
+    ],
+)
 
 
 def record_agent_invocation_duration(
@@ -148,6 +186,18 @@ def record_workflow_invocation_duration(
   _workflow_invocation_duration.record(elapsed_s, attributes=attrs)
 
 
+def record_invoke_agent_inference_calls(agent_name: str, count: int) -> None:
+  """Records the number of inference (model) calls in an agent invocation."""
+  attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name}
+  _invoke_agent_inference_calls.record(count, attributes=attrs)
+
+
+def record_invoke_agent_tool_calls(agent_name: str, count: int) -> None:
+  """Records the number of tool calls in an agent invocation."""
+  attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name}
+  _invoke_agent_tool_calls.record(count, attributes=attrs)
+
+
 def record_agent_workflow_steps(agent_name: str, events: list[Event]):
   """Records the number of steps in the agent workflow by counting the number of events."""
   attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name}

src/google/adk/telemetry/node_tracing.py

@@ -82,10 +82,12 @@ async def start_as_current_node_span(
   - `invoke_workflow {workflow.name}`
   - `invoke_node {node.name}`
 
-  invoke_agent spans align with OpenTelemetry Semantic Conventions (semconv) version 1.36 spans for backwards compatibility.
+  invoke_agent spans align with OpenTelemetry Semantic Conventions (semconv)
+  version 1.36 spans for backwards compatibility.
   https://github.com/open-telemetry/semantic-conventions/blob/v1.36.0/docs/gen-ai/README.md
 
-  invoke_workflow spans align with semconv version 1.41, because these were not included in any prior releases.
+  invoke_workflow spans align with semconv version 1.41, because these were not
+  included in any prior releases.
   https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/README.md
 
   invoke_node spans are not present in any semconv release.

tests/unittests/telemetry/functional_node_test_cases.py

@@ -2621,6 +2621,12 @@
             value=NON_DETERMINISTIC,
         ),
     }),
+    "gen_ai.invoke_agent.inference_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2),
+    }),
+    "gen_ai.invoke_agent.tool_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1),
+    }),
 }
 
 
@@ -2665,6 +2671,12 @@
             value=NON_DETERMINISTIC,
         ),
     }),
+    "gen_ai.invoke_agent.inference_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2),
+    }),
+    "gen_ai.invoke_agent.tool_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1),
+    }),
 }
 
 

tests/unittests/telemetry/functional_test_cases.py

@@ -2306,6 +2306,12 @@
             value=NON_DETERMINISTIC,
         ),
     }),
+    "gen_ai.invoke_agent.inference_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2),
+    }),
+    "gen_ai.invoke_agent.tool_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1),
+    }),
 }
 
 
@@ -2350,6 +2356,12 @@
             value=NON_DETERMINISTIC,
         ),
     }),
+    "gen_ai.invoke_agent.inference_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2),
+    }),
+    "gen_ai.invoke_agent.tool_calls": frozenset({
+        MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1),
+    }),
 }
 
 

tests/unittests/telemetry/functional_test_helpers.py

@@ -334,6 +334,16 @@ class HistogramSpec(NamedTuple):
         attr="_workflow_invocation_duration",
         metric_name="gen_ai.invoke_workflow.duration",
     ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_invoke_agent_inference_calls",
+        metric_name="gen_ai.invoke_agent.inference_calls",
+    ),
+    HistogramSpec(
+        module=_metrics,
+        attr="_invoke_agent_tool_calls",
+        metric_name="gen_ai.invoke_agent.tool_calls",
+    ),
 )