Merge branch 'main' into marlzrana/fix-nullable-types-w-gen-ai-bump

MarlzRana · web-flow · commit ad28b82cbdbb · 2025-10-27T08:41:33.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,43 +25,43 @@ classifiers = [ # List of https://pypi.org/classifiers/
 ]
 dependencies = [
   # go/keep-sorted start
-  "PyYAML>=6.0.2, <7.0.0",                                  # For APIHubToolset.
-  "absolufy-imports>=0.3.1, <1.0.0",                        # For Agent Engine deployment.
-  "anyio>=4.9.0, <5.0.0;python_version>='3.10'",            # For MCP Session Manager
-  "authlib>=1.5.1, <2.0.0",                                 # For RestAPI Tool
-  "click>=8.1.8, <9.0.0",                                   # For CLI tools
-  "fastapi>=0.115.0, <1.0.0",                               # FastAPI framework
-  "google-api-python-client>=2.157.0, <3.0.0",              # Google API client discovery
-  "google-cloud-aiplatform[agent_engines]>=1.121.0, <2.0.0",# For VertexAI integrations, e.g. example store.
-  "google-cloud-bigtable>=2.32.0",                          # For Bigtable database
-  "google-cloud-discoveryengine>=0.13.12, <0.14.0",         # For Discovery Engine Search Tool
-  "google-cloud-secret-manager>=2.22.0, <3.0.0",            # Fetching secrets in RestAPI Tool
-  "google-cloud-spanner>=3.56.0, <4.0.0",                   # For Spanner database
-  "google-cloud-speech>=2.30.0, <3.0.0",                    # For Audio Transcription
-  "google-cloud-storage>=2.18.0, <3.0.0",                   # For GCS Artifact service
-  "google-genai>=1.45.0, <2.0.0",                                   # Google GenAI SDK
-  "graphviz>=0.20.2, <1.0.0",                               # Graphviz for graph rendering
-  "mcp>=1.8.0, <2.0.0;python_version>='3.10'",              # For MCP Toolset
-  "opentelemetry-api>=1.37.0, <=1.37.0",                    # OpenTelemetry - limit upper version for sdk and api to not risk breaking changes from unstable _logs package.
+  "PyYAML>=6.0.2, <7.0.0",                                   # For APIHubToolset.
+  "absolufy-imports>=0.3.1, <1.0.0",                         # For Agent Engine deployment.
+  "anyio>=4.9.0, <5.0.0;python_version>='3.10'",             # For MCP Session Manager
+  "authlib>=1.5.1, <2.0.0",                                  # For RestAPI Tool
+  "click>=8.1.8, <9.0.0",                                    # For CLI tools
+  "fastapi>=0.115.0, <1.119.0",                              # FastAPI framework
+  "google-api-python-client>=2.157.0, <3.0.0",               # Google API client discovery
+  "google-cloud-aiplatform[agent_engines]>=1.121.0, <2.0.0", # For VertexAI integrations, e.g. example store.
+  "google-cloud-bigtable>=2.32.0",                           # For Bigtable database
+  "google-cloud-discoveryengine>=0.13.12, <0.14.0",          # For Discovery Engine Search Tool
+  "google-cloud-secret-manager>=2.22.0, <3.0.0",             # Fetching secrets in RestAPI Tool
+  "google-cloud-spanner>=3.56.0, <4.0.0",                    # For Spanner database
+  "google-cloud-speech>=2.30.0, <3.0.0",                     # For Audio Transcription
+  "google-cloud-storage>=2.18.0, <3.0.0",                    # For GCS Artifact service
+  "google-genai>=1.45.0, <2.0.0",                            # Google GenAI SDK
+  "graphviz>=0.20.2, <1.0.0",                                # Graphviz for graph rendering
+  "mcp>=1.8.0, <2.0.0;python_version>='3.10'",               # For MCP Toolset
+  "opentelemetry-api>=1.37.0, <=1.37.0",                     # OpenTelemetry - limit upper version for sdk and api to not risk breaking changes from unstable _logs package.
   "opentelemetry-exporter-gcp-logging>=1.9.0a0, <2.0.0",
   "opentelemetry-exporter-gcp-monitoring>=1.9.0a0, <2.0.0",
   "opentelemetry-exporter-gcp-trace>=1.9.0, <2.0.0",
   "opentelemetry-exporter-otlp-proto-http>=1.36.0",
   "opentelemetry-resourcedetector-gcp>=1.9.0a0, <2.0.0",
   "opentelemetry-sdk>=1.37.0, <=1.37.0",
-  "pydantic>=2.0, <3.0.0",                                  # For data validation/models
-  "python-dateutil>=2.9.0.post0, <3.0.0",                   # For Vertext AI Session Service
-  "python-dotenv>=1.0.0, <2.0.0",                           # To manage environment variables
+  "pydantic>=2.0, <3.0.0",                                   # For data validation/models
+  "python-dateutil>=2.9.0.post0, <3.0.0",                    # For Vertext AI Session Service
+  "python-dotenv>=1.0.0, <2.0.0",                            # To manage environment variables
   "requests>=2.32.4, <3.0.0",
-  "sqlalchemy-spanner>=1.14.0",                             # Spanner database session service
-  "sqlalchemy>=2.0, <3.0.0",                                # SQL database ORM
-  "starlette>=0.46.2, <1.0.0",                              # For FastAPI CLI
-  "tenacity>=8.0.0, <9.0.0",                                # For Retry management
+  "sqlalchemy-spanner>=1.14.0",                              # Spanner database session service
+  "sqlalchemy>=2.0, <3.0.0",                                 # SQL database ORM
+  "starlette>=0.46.2, <1.0.0",                               # For FastAPI CLI
+  "tenacity>=8.0.0, <9.0.0",                                 # For Retry management
   "typing-extensions>=4.5, <5",
-  "tzlocal>=5.3, <6.0",                                     # Time zone utilities
-  "uvicorn>=0.34.0, <1.0.0",                                # ASGI server for FastAPI
-  "watchdog>=6.0.0, <7.0.0",                                # For file change detection and hot reload
-  "websockets>=15.0.1, <16.0.0",                            # For BaseLlmFlow
+  "tzlocal>=5.3, <6.0",                                      # Time zone utilities
+  "uvicorn>=0.34.0, <1.0.0",                                 # ASGI server for FastAPI
+  "watchdog>=6.0.0, <7.0.0",                                 # For file change detection and hot reload
+  "websockets>=15.0.1, <16.0.0",                             # For BaseLlmFlow
   # go/keep-sorted end
 ]
 dynamic = ["version"]
@@ -111,13 +111,13 @@ eval = [
 test = [
   # go/keep-sorted start
   "a2a-sdk>=0.3.0,<0.4.0;python_version>='3.10'",
-  "anthropic>=0.43.0",                      # For anthropic model tests
-  "kubernetes>=29.0.0",                     # For GkeCodeExecutor
+  "anthropic>=0.43.0",                            # For anthropic model tests
+  "kubernetes>=29.0.0",                           # For GkeCodeExecutor
   "langchain-community>=0.3.17",
-  "langgraph>=0.2.60, <0.4.8",              # For LangGraphAgent
-  "litellm>=1.75.5, <2.0.0",                # For LiteLLM tests
-  "llama-index-readers-file>=0.4.0",        # For retrieval tests
-  "openai>=1.100.2",                        # For LiteLLM
+  "langgraph>=0.2.60, <0.4.8",                    # For LangGraphAgent
+  "litellm>=1.75.5, <2.0.0",                      # For LiteLLM tests
+  "llama-index-readers-file>=0.4.0",              # For retrieval tests
+  "openai>=1.100.2",                              # For LiteLLM
   "pytest>=8.3.4",
   "pytest-asyncio>=0.25.0",
   "pytest-mock>=3.14.0",
@@ -139,22 +139,20 @@ docs = [
 
 # Optional extensions
 extensions = [
-  "anthropic>=0.43.0",                         # For anthropic model support
-  "beautifulsoup4>=3.2.2",                     # For load_web_page tool.
-  "crewai[tools];python_version>='3.10'",      # For CrewaiTool
-  "docker>=7.0.0",                             # For ContainerCodeExecutor
-  "kubernetes>=29.0.0",                        # For GkeCodeExecutor
-  "langgraph>=0.2.60, <0.4.8",                 # For LangGraphAgent
-  "litellm>=1.75.5",                           # For LiteLlm class. Currently has OpenAI limitations. TODO: once LiteLlm fix it
-  "llama-index-readers-file>=0.4.0",           # For retrieval using LlamaIndex.
-  "llama-index-embeddings-google-genai>=0.3.0",# For files retrieval using LlamaIndex.
-  "lxml>=5.3.0",                               # For load_web_page tool.
-  "toolbox-core>=0.1.0",                       # For tools.toolbox_toolset.ToolboxToolset
+  "anthropic>=0.43.0",                          # For anthropic model support
+  "beautifulsoup4>=3.2.2",                      # For load_web_page tool.
+  "crewai[tools];python_version>='3.10'",       # For CrewaiTool
+  "docker>=7.0.0",                              # For ContainerCodeExecutor
+  "kubernetes>=29.0.0",                         # For GkeCodeExecutor
+  "langgraph>=0.2.60, <0.4.8",                  # For LangGraphAgent
+  "litellm>=1.75.5",                            # For LiteLlm class. Currently has OpenAI limitations. TODO: once LiteLlm fix it
+  "llama-index-readers-file>=0.4.0",            # For retrieval using LlamaIndex.
+  "llama-index-embeddings-google-genai>=0.3.0", # For files retrieval using LlamaIndex.
+  "lxml>=5.3.0",                                # For load_web_page tool.
+  "toolbox-core>=0.1.0",                        # For tools.toolbox_toolset.ToolboxToolset
 ]
 
-otel-gcp = [
-  "opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0",
-]
+otel-gcp = ["opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0"]
 
 
 [tool.pyink]
diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py
@@ -155,8 +155,8 @@ class MyAgent(BaseAgent):
 
   Returns:
     Optional[types.Content]: The content to return to the user.
-      When the content is present, the provided content will be used as agent
-      response and appended to event history as agent response.
+      When the content is present, an additional event with the provided content
+      will be appended to event history as an additional agent response.
   """
 
   def _load_agent_state(
diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py
@@ -305,7 +305,9 @@ async def _execute_single_function_call_async(
     else:
       raise tool_error
 
-  with tracer.start_as_current_span(f'execute_tool {tool.name}'):
+  async def _run_with_trace():
+    nonlocal function_args
+
     # Step 1: Check if plugin before_tool_callback overrides the function
     # response.
     function_response = (
@@ -391,13 +393,23 @@ async def _execute_single_function_call_async(
     function_response_event = __build_response_event(
         tool, function_response, tool_context, invocation_context
     )
-    trace_tool_call(
-        tool=tool,
-        args=function_args,
-        function_response_event=function_response_event,
-    )
     return function_response_event
 
+  with tracer.start_as_current_span(f'execute_tool {tool.name}'):
+    try:
+      function_response_event = await _run_with_trace()
+      trace_tool_call(
+          tool=tool,
+          args=function_args,
+          function_response_event=function_response_event,
+      )
+      return function_response_event
+    except:
+      trace_tool_call(
+          tool=tool, args=function_args, function_response_event=None
+      )
+      raise
+
 
 async def handle_function_calls_live(
     invocation_context: InvocationContext,
@@ -467,13 +479,17 @@ async def _execute_single_function_call_live(
   tool, tool_context = _get_tool_and_context(
       invocation_context, function_call, tools_dict
   )
-  with tracer.start_as_current_span(f'execute_tool {tool.name}'):
+
+  function_args = (
+      copy.deepcopy(function_call.args) if function_call.args else {}
+  )
+
+  async def _run_with_trace():
+    nonlocal function_args
+
     # Do not use "args" as the variable name, because it is a reserved keyword
     # in python debugger.
     # Make a deep copy to avoid being modified.
-    function_args = (
-        copy.deepcopy(function_call.args) if function_call.args else {}
-    )
     function_response = None
 
     # Handle before_tool_callbacks - iterate through the canonical callback
@@ -527,13 +543,23 @@ async def _execute_single_function_call_live(
     function_response_event = __build_response_event(
         tool, function_response, tool_context, invocation_context
     )
-    trace_tool_call(
-        tool=tool,
-        args=function_args,
-        function_response_event=function_response_event,
-    )
     return function_response_event
 
+  with tracer.start_as_current_span(f'execute_tool {tool.name}'):
+    try:
+      function_response_event = await _run_with_trace()
+      trace_tool_call(
+          tool=tool,
+          args=function_args,
+          function_response_event=function_response_event,
+      )
+      return function_response_event
+    except:
+      trace_tool_call(
+          tool=tool, args=function_args, function_response_event=None
+      )
+      raise
+
 
 async def _process_function_live_helper(
     tool,
diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py
@@ -26,6 +26,7 @@
 import json
 import os
 from typing import Any
+from typing import Optional
 from typing import TYPE_CHECKING
 
 from google.genai import types
@@ -118,7 +119,7 @@ def trace_agent_invocation(
 def trace_tool_call(
     tool: BaseTool,
     args: dict[str, Any],
-    function_response_event: Event,
+    function_response_event: Optional[Event],
 ):
   """Traces tool call.
 
@@ -154,7 +155,8 @@ def trace_tool_call(
   tool_call_id = '<not specified>'
   tool_response = '<not specified>'
   if (
-      function_response_event.content is not None
+      function_response_event is not None
+      and function_response_event.content is not None
       and function_response_event.content.parts
   ):
     response_parts = function_response_event.content.parts
@@ -169,7 +171,8 @@ def trace_tool_call(
 
   if not isinstance(tool_response, dict):
     tool_response = {'result': tool_response}
-  span.set_attribute('gcp.vertex.agent.event_id', function_response_event.id)
+  if function_response_event is not None:
+    span.set_attribute('gcp.vertex.agent.event_id', function_response_event.id)
   if _should_add_request_response_to_spans():
     span.set_attribute(
         'gcp.vertex.agent.tool_response',
diff --git a/tests/unittests/telemetry/test_functional.py b/tests/unittests/telemetry/test_functional.py
@@ -12,18 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import gc
 import sys
-from unittest import mock
 
 from google.adk.agents import base_agent
 from google.adk.agents.llm_agent import Agent
 from google.adk.models.base_llm import BaseLlm
+from google.adk.models.llm_response import LlmResponse
 from google.adk.telemetry import tracing
 from google.adk.tools import FunctionTool
 from google.adk.utils.context_utils import Aclosing
+from google.genai.types import Content
 from google.genai.types import Part
-from opentelemetry.version import __version__
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 import pytest
 
 from ..testing_utils import MockModel
@@ -63,27 +67,27 @@ async def test_runner(test_agent: Agent) -> TestInMemoryRunner:
 
 
 @pytest.fixture
-def mock_start_as_current_span(monkeypatch: pytest.MonkeyPatch) -> mock.Mock:
-  mock_context_manager = mock.MagicMock()
-  mock_context_manager.__enter__.return_value = mock.Mock()
-  mock_start_as_current_span = mock.Mock()
-  mock_start_as_current_span.return_value = mock_context_manager
+def span_exporter(monkeypatch: pytest.MonkeyPatch) -> InMemorySpanExporter:
+  tracer_provider = TracerProvider()
+  span_exporter = InMemorySpanExporter()
+  tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
+  real_tracer = tracer_provider.get_tracer(__name__)
 
   def do_replace(tracer):
     monkeypatch.setattr(
-        tracer, 'start_as_current_span', mock_start_as_current_span
+        tracer, 'start_as_current_span', real_tracer.start_as_current_span
     )
 
   do_replace(tracing.tracer)
   do_replace(base_agent.tracer)
 
-  return mock_start_as_current_span
+  return span_exporter
 
 
 @pytest.mark.asyncio
 async def test_tracer_start_as_current_span(
     test_runner: TestInMemoryRunner,
-    mock_start_as_current_span: mock.Mock,
+    span_exporter: InMemorySpanExporter,
 ):
   """Test creation of multiple spans in an E2E runner invocation.
 
@@ -112,18 +116,49 @@ def wrapped_firstiter(coro):
       pass
 
   # Assert
-  expected_start_as_current_span_calls = [
-      mock.call('invocation'),
-      mock.call('execute_tool some_tool'),
-      mock.call('invoke_agent some_root_agent'),
-      mock.call('call_llm'),
-      mock.call('call_llm'),
+  spans = span_exporter.get_finished_spans()
+  assert list(sorted(span.name for span in spans)) == [
+      'call_llm',
+      'call_llm',
+      'execute_tool some_tool',
+      'invocation',
+      'invoke_agent some_root_agent',
   ]
 
-  mock_start_as_current_span.assert_has_calls(
-      expected_start_as_current_span_calls,
-      any_order=True,
+
+@pytest.mark.asyncio
+async def test_exception_preserves_attributes(
+    test_model: BaseLlm, span_exporter: InMemorySpanExporter
+):
+  """Test when an exception occurs during tool execution, span attributes are still present on spans where they are expected."""
+
+  # Arrange
+  async def some_tool():
+    raise ValueError('This tool always fails')
+
+  test_agent = Agent(
+      name='some_root_agent',
+      model=test_model,
+      tools=[
+          FunctionTool(some_tool),
+      ],
   )
-  assert mock_start_as_current_span.call_count == len(
-      expected_start_as_current_span_calls
+
+  test_runner = TestInMemoryRunner(test_agent)
+
+  # Act
+  with pytest.raises(ValueError, match='This tool always fails'):
+    async with Aclosing(
+        test_runner.run_async_with_new_session_agen('')
+    ) as agen:
+      async for _ in agen:
+        pass
+
+  # Assert
+  spans = span_exporter.get_finished_spans()
+  assert len(spans) > 1
+  assert all(
+      span.attributes is not None and len(span.attributes) > 0
+      for span in spans
+      if span.name != 'invocation'  # not expected to have attributes
   )
