Feature/disable telemetry

src/google/adk/agents/base_agent.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import inspect
 import logging
 from typing import Any
@@ -44,6 +45,7 @@
 from ..telemetry.tracing import tracer
 from ..utils.context_utils import Aclosing
 from ..utils.feature_decorator import experimental
+from ..utils.telemetry_utils import is_telemetry_enabled
 from .base_agent_config import BaseAgentConfig
 from .callback_context import CallbackContext
 
@@ -132,6 +134,8 @@ class MyAgent(BaseAgent):
   """
   sub_agents: list[BaseAgent] = Field(default_factory=list)
   """The sub-agents of this agent."""
+  disable_telemetry: bool = False
+  """Whether to disable telemetry for this agent."""
 
   before_agent_callback: Optional[BeforeAgentCallback] = None
   """Callback or list of callbacks to be invoked before the agent run.
@@ -282,24 +286,20 @@ async def run_async(
       Event: the events generated by the agent.
     """
 
-    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
-      ctx = self._create_invocation_context(parent_context)
-      tracing.trace_agent_invocation(span, self, ctx)
-      if event := await self._handle_before_agent_callback(ctx):
-        yield event
-      if ctx.end_invocation:
-        return
+    span_context = contextlib.nullcontext()
+    if is_telemetry_enabled(self):
+      span_context = tracer.start_as_current_span(f'invoke_agent {self.name}')
 
-      async with Aclosing(self._run_async_impl(ctx)) as agen:
+    with span_context as span:
+      ctx = self._create_invocation_context(parent_context)
+      if span:
+        tracing.trace_agent_invocation(span, self, ctx)
+      async with Aclosing(
+          self._run_callbacks_and_impl(ctx, mode='async')
+      ) as agen:
         async for event in agen:
           yield event
 
-      if ctx.end_invocation:
-        return
-
-      if event := await self._handle_after_agent_callback(ctx):
-        yield event
-
   @final
   async def run_live(
       self,
@@ -315,19 +315,15 @@ async def run_live(
       Event: the events generated by the agent.
     """
 
-    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
-      ctx = self._create_invocation_context(parent_context)
-      tracing.trace_agent_invocation(span, self, ctx)
-      if event := await self._handle_before_agent_callback(ctx):
-        yield event
-      if ctx.end_invocation:
-        return
-
-      async with Aclosing(self._run_live_impl(ctx)) as agen:
-        async for event in agen:
-          yield event
+    span_context = contextlib.nullcontext()
+    if is_telemetry_enabled(self):
+      span_context = tracer.start_as_current_span(f'invoke_agent {self.name}')
 
-      if event := await self._handle_after_agent_callback(ctx):
+    with span_context as span:
+      ctx = self._create_invocation_context(parent_context)
+      if span:
+        tracing.trace_agent_invocation(span, self, ctx)
+      async for event in self._run_callbacks_and_impl(ctx, mode='live'):
         yield event
 
   async def _run_async_impl(
@@ -362,6 +358,39 @@ async def _run_live_impl(
     )
     yield  # AsyncGenerator requires having at least one yield statement
 
+  async def _run_callbacks_and_impl(
+      self, ctx: InvocationContext, mode: str = 'async'
+  ) -> AsyncGenerator[Event, None]:
+    """Runs the before and after agent callbacks around the core agent logic.
+    Args:
+      ctx: InvocationContext, the invocation context for this agent.
+      mode: str, either 'async' or 'live', indicating which core agent logic to run.
+    Yields:
+      Event: the events generated by the agent.
+    """
+    if event := await self._handle_before_agent_callback(ctx):
+      yield event
+    if ctx.end_invocation:
+      return
+    if mode.lower() == 'async':
+      async with Aclosing(self._run_async_impl(ctx)) as agen:
+        async for event in agen:
+          yield event
+    elif mode.lower() == 'live':
+      async with Aclosing(self._run_live_impl(ctx)) as agen:
+        async for event in agen:
+          yield event
+    else:
+      raise ValueError(
+          f"Invalid mode: {mode}. Must be either 'async' or 'live'."
+      )
+
+    if ctx.end_invocation:
+      return
+
+    if event := await self._handle_after_agent_callback(ctx):
+      yield event
+
   @property
   def root_agent(self) -> BaseAgent:
     """Gets the root agent of this agent."""

src/google/adk/agents/llm_agent.py

@@ -55,6 +55,7 @@
 from ..tools.tool_context import ToolContext
 from ..utils.context_utils import Aclosing
 from ..utils.feature_decorator import experimental
+from ..utils.telemetry_utils import is_telemetry_enabled
 from .base_agent import BaseAgent
 from .base_agent import BaseAgentState
 from .base_agent_config import BaseAgentConfig
@@ -814,6 +815,10 @@ def __maybe_save_output_to_state(self, event: Event):
 
   @model_validator(mode='after')
   def __model_validator_after(self) -> LlmAgent:
+    root_agent = getattr(self, 'root_agent', None) or self
+    disable_telemetry: bool = not is_telemetry_enabled(root_agent)
+    if hasattr(self.model, 'disable_telemetry'):
+      self.model.disable_telemetry = disable_telemetry
     return self
 
   @field_validator('generate_content_config', mode='after')

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -16,6 +16,7 @@
 
 from abc import ABC
 import asyncio
+import contextlib
 import datetime
 import inspect
 import logging
@@ -50,6 +51,7 @@
 from ...tools.google_search_tool import google_search
 from ...tools.tool_context import ToolContext
 from ...utils.context_utils import Aclosing
+from ...utils.telemetry_utils import is_telemetry_enabled
 from .audio_cache_manager import AudioCacheManager
 
 if TYPE_CHECKING:
@@ -129,13 +131,16 @@ async def run_live(
         async with llm.connect(llm_request) as llm_connection:
           if llm_request.contents:
             # Sends the conversation history to the model.
-            with tracer.start_as_current_span('send_data'):
-              # Combine regular contents with audio/transcription from session
+            span_context = contextlib.nullcontext()
+            if is_telemetry_enabled(invocation_context.agent):
+              span_context = tracer.start_as_current_span('send_data')
+            with span_context as span:
               logger.debug('Sending history to model: %s', llm_request.contents)
               await llm_connection.send_history(llm_request.contents)
-              trace_send_data(
-                  invocation_context, event_id, llm_request.contents
-              )
+              if span:
+                trace_send_data(
+                    invocation_context, event_id, llm_request.contents
+                )
 
           send_task = asyncio.create_task(
               self._send_to_model(llm_connection, invocation_context)
@@ -752,7 +757,10 @@ async def _call_llm_async(
     llm = self.__get_llm(invocation_context)
 
     async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
-      with tracer.start_as_current_span('call_llm'):
+      span_context = contextlib.nullcontext()
+      if is_telemetry_enabled(invocation_context.agent):
+        span_context = tracer.start_as_current_span('call_llm')
+      with span_context:
         if invocation_context.run_config.support_cfc:
           invocation_context.live_request_queue = LiveRequestQueue()
           responses_generator = self.run_live(invocation_context)

src/google/adk/flows/llm_flows/functions.py

@@ -42,6 +42,7 @@
 from ...tools.tool_confirmation import ToolConfirmation
 from ...tools.tool_context import ToolContext
 from ...utils.context_utils import Aclosing
+from ...utils.telemetry_utils import is_telemetry_enabled
 
 if TYPE_CHECKING:
   from ...agents.llm_agent import LlmAgent
@@ -255,7 +256,7 @@ async def handle_function_call_list_async(
       function_response_events
   )
 
-  if len(function_response_events) > 1:
+  if len(function_response_events) > 1 and is_telemetry_enabled(agent):
     # this is needed for debug traces of parallel calls
     # individual response with tool.name is traced in __build_response_event
     # (we drop tool.name from span name here as this is merged event)
@@ -425,6 +426,8 @@ async def _run_with_trace():
     )
     return function_response_event
 
+  if not is_telemetry_enabled(agent):
+    return await _run_with_trace()
   with tracer.start_as_current_span(f'execute_tool {tool.name}'):
     try:
       function_response_event = await _run_with_trace()
@@ -486,7 +489,7 @@ async def handle_function_calls_live(
   merged_event = merge_parallel_function_response_events(
       function_response_events
   )
-  if len(function_response_events) > 1:
+  if len(function_response_events) > 1 and is_telemetry_enabled(agent):
     # this is needed for debug traces of parallel calls
     # individual response with tool.name is traced in __build_response_event
     # (we drop tool.name from span name here as this is merged event)
@@ -575,6 +578,8 @@ async def _run_with_trace():
     )
     return function_response_event
 
+  if not is_telemetry_enabled(agent):
+    return await _run_with_trace()
   with tracer.start_as_current_span(f'execute_tool {tool.name}'):
     try:
       function_response_event = await _run_with_trace()

src/google/adk/models/gemini_context_cache_manager.py

@@ -15,7 +15,7 @@
 """Manages context cache lifecycle for Gemini models."""
 
 from __future__ import annotations
-
+import contextlib
 import hashlib
 import json
 import logging
@@ -24,6 +24,7 @@
 from typing import TYPE_CHECKING
 
 from google.genai import types
+from opentelemetry.trace import Span
 
 from ..utils.feature_decorator import experimental
 from .cache_metadata import CacheMetadata
@@ -45,13 +46,15 @@ class GeminiContextCacheManager:
   cache compatibility and implements efficient caching strategies.
   """
 
-  def __init__(self, genai_client: Client):
+  def __init__(self, genai_client: Client, disable_telemetry: bool = False):
     """Initialize cache manager with shared client.
 
     Args:
         genai_client: The GenAI client to use for cache operations.
+        disable_telemetry: A bool to flag whether or not telemetry should be disabled.
     """
     self.genai_client = genai_client
+    self.disable_telemetry = disable_telemetry
 
   async def handle_context_caching(
       self, llm_request: LlmRequest
@@ -356,9 +359,13 @@ async def _create_gemini_cache(
     Returns:
         Cache metadata with precise creation timestamp
     """
-    from ..telemetry.tracing import tracer
 
-    with tracer.start_as_current_span("create_cache") as span:
+    span_context = contextlib.nullcontext()
+    if not self.disable_telemetry:
+      from ..telemetry.tracing import tracer
+      span_context = tracer.start_as_current_span("create_cache")
+
+    with span_context as span:
       # Prepare cache contents (first N contents + system instruction + tools)
       cache_contents = llm_request.contents[:cache_contents_count]
 
@@ -386,9 +393,10 @@ async def _create_gemini_cache(
       if llm_request.config and llm_request.config.tool_config:
         cache_config.tool_config = llm_request.config.tool_config
 
-      span.set_attribute("cache_contents_count", cache_contents_count)
-      span.set_attribute("model", llm_request.model)
-      span.set_attribute("ttl_seconds", llm_request.cache_config.ttl_seconds)
+      if span is not None:
+        span.set_attribute("cache_contents_count", cache_contents_count)
+        span.set_attribute("model", llm_request.model)
+        span.set_attribute("ttl_seconds", llm_request.cache_config.ttl_seconds)
 
       logger.debug(
           "Creating cache with model %s and config: %s",
@@ -403,7 +411,8 @@ async def _create_gemini_cache(
       created_at = time.time()
       logger.info("Cache created successfully: %s", cached_content.name)
 
-      span.set_attribute("cache_name", cached_content.name)
+      if span is not None:
+        span.set_attribute("cache_name", cached_content.name)
 
       # Return complete cache metadata with precise timing
       return CacheMetadata(

src/google/adk/models/google_llm.py

@@ -127,6 +127,11 @@ class Gemini(BaseLlm):
   ```
   """
 
+  disable_telemetry: bool = False
+  """A bool to flag whether or not telemetry should be being disabled for
+  Gemini LLM interactions.
+  """
+
   @classmethod
   @override
   def supported_models(cls) -> list[str]:
@@ -165,18 +170,29 @@ async def generate_content_async(
     cache_metadata = None
     cache_manager = None
     if llm_request.cache_config:
-      from ..telemetry.tracing import tracer
       from .gemini_context_cache_manager import GeminiContextCacheManager
 
-      with tracer.start_as_current_span('handle_context_caching') as span:
-        cache_manager = GeminiContextCacheManager(self.api_client)
+      if not self.disable_telemetry:
+        from ..telemetry.tracing import tracer
+
+        with tracer.start_as_current_span('handle_context_caching') as span:
+          cache_manager = GeminiContextCacheManager(
+              self.api_client, disable_telemetry=self.disable_telemetry
+          )
+          cache_metadata = await cache_manager.handle_context_caching(
+              llm_request
+          )
+          if cache_metadata:
+            if cache_metadata.cache_name:
+              span.set_attribute('cache_action', 'active_cache')
+              span.set_attribute('cache_name', cache_metadata.cache_name)
+            else:
+              span.set_attribute('cache_action', 'fingerprint_only')
+      else:
+        cache_manager = GeminiContextCacheManager(
+            self.api_client, disable_telemetry=self.disable_telemetry
+        )
         cache_metadata = await cache_manager.handle_context_caching(llm_request)
-        if cache_metadata:
-          if cache_metadata.cache_name:
-            span.set_attribute('cache_action', 'active_cache')
-            span.set_attribute('cache_name', cache_metadata.cache_name)
-          else:
-            span.set_attribute('cache_action', 'fingerprint_only')
 
     logger.info(
         'Sending out request, model: %s, backend: %s, stream: %s',