google-adk 1.1.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

google/adk/flows/llm_flows/base_llm_flow.py

@@ -23,6 +23,7 @@ from typing import cast
 from typing import Optional
 from typing import TYPE_CHECKING
 
+from google.genai import types
 from websockets.exceptions import ConnectionClosedOK
 
 from . import functions
@@ -50,6 +51,8 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger('google_adk.' + __name__)
 
+_ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'
+
 
 class BaseLlmFlow(ABC):
   """A basic flow that calls the LLM in a loop until a final response is generated.
@@ -281,6 +284,12 @@ class BaseLlmFlow(ABC):
         yield event
       if not last_event or last_event.is_final_response():
         break
+      if last_event.partial:
+        # TODO: handle this in BaseLlm level.
+        raise ValueError(
+            f"Last event shouldn't be partial. LLM max output limit may be"
+            f' reached.'
+        )
 
   async def _run_one_step_async(
       self,
@@ -493,6 +502,16 @@ class BaseLlmFlow(ABC):
       yield response
       return
 
+    llm_request.config = llm_request.config or types.GenerateContentConfig()
+    llm_request.config.labels = llm_request.config.labels or {}
+
+    # Add agent name as a label to the llm_request. This will help with slicing
+    # the billing reports on a per-agent basis.
+    if _ADK_AGENT_NAME_LABEL_KEY not in llm_request.config.labels:
+      llm_request.config.labels[_ADK_AGENT_NAME_LABEL_KEY] = (
+          invocation_context.agent.name
+      )
+
     # Calls the LLM.
     llm = self.__get_llm(invocation_context)
     with tracer.start_as_current_span('call_llm'):

google/adk/flows/llm_flows/contents.py

@@ -170,10 +170,10 @@ def _rearrange_events_for_latest_function_response(
   for idx in range(function_call_event_idx + 1, len(events) - 1):
     event = events[idx]
     function_responses = event.get_function_responses()
-    if (
-        function_responses
-        and function_responses[0].id in function_responses_ids
-    ):
+    if function_responses and any([
+        function_response.id in function_responses_ids
+        for function_response in function_responses
+    ]):
       function_response_events.append(event)
   function_response_events.append(events[-1])
 

google/adk/flows/llm_flows/functions.py

@@ -32,8 +32,8 @@ from ...agents.invocation_context import InvocationContext
 from ...auth.auth_tool import AuthToolArguments
 from ...events.event import Event
 from ...events.event_actions import EventActions
+from ...telemetry import trace_merged_tool_calls
 from ...telemetry import trace_tool_call
-from ...telemetry import trace_tool_response
 from ...telemetry import tracer
 from ...tools.base_tool import BaseTool
 from ...tools.tool_context import ToolContext
@@ -148,62 +148,69 @@ async def handle_function_calls_async(
         function_call,
         tools_dict,
     )
-    # do not use "args" as the variable name, because it is a reserved keyword
-    # in python debugger.
-    function_args = function_call.args or {}
-    function_response: Optional[dict] = None
-
-    for callback in agent.canonical_before_tool_callbacks:
-      function_response = callback(
-          tool=tool, args=function_args, tool_context=tool_context
-      )
-      if inspect.isawaitable(function_response):
-        function_response = await function_response
-      if function_response:
-        break
-
-    if not function_response:
-      function_response = await __call_tool_async(
-          tool, args=function_args, tool_context=tool_context
-      )
 
-    for callback in agent.canonical_after_tool_callbacks:
-      altered_function_response = callback(
+    with tracer.start_as_current_span(f'execute_tool {tool.name}'):
+      # do not use "args" as the variable name, because it is a reserved keyword
+      # in python debugger.
+      function_args = function_call.args or {}
+      function_response: Optional[dict] = None
+
+      for callback in agent.canonical_before_tool_callbacks:
+        function_response = callback(
+            tool=tool, args=function_args, tool_context=tool_context
+        )
+        if inspect.isawaitable(function_response):
+          function_response = await function_response
+        if function_response:
+          break
+
+      if not function_response:
+        function_response = await __call_tool_async(
+            tool, args=function_args, tool_context=tool_context
+        )
+
+      for callback in agent.canonical_after_tool_callbacks:
+        altered_function_response = callback(
+            tool=tool,
+            args=function_args,
+            tool_context=tool_context,
+            tool_response=function_response,
+        )
+        if inspect.isawaitable(altered_function_response):
+          altered_function_response = await altered_function_response
+        if altered_function_response is not None:
+          function_response = altered_function_response
+          break
+
+      if tool.is_long_running:
+        # Allow long running function to return None to not provide function response.
+        if not function_response:
+          continue
+
+      # Builds the function response event.
+      function_response_event = __build_response_event(
+          tool, function_response, tool_context, invocation_context
+      )
+      trace_tool_call(
           tool=tool,
           args=function_args,
-          tool_context=tool_context,
-          tool_response=function_response,
+          function_response_event=function_response_event,
       )
-      if inspect.isawaitable(altered_function_response):
-        altered_function_response = await altered_function_response
-      if altered_function_response is not None:
-        function_response = altered_function_response
-        break
-
-    if tool.is_long_running:
-      # Allow long running function to return None to not provide function response.
-      if not function_response:
-        continue
-
-    # Builds the function response event.
-    function_response_event = __build_response_event(
-        tool, function_response, tool_context, invocation_context
-    )
-    function_response_events.append(function_response_event)
+      function_response_events.append(function_response_event)
 
   if not function_response_events:
     return None
   merged_event = merge_parallel_function_response_events(
       function_response_events
   )
+
   if len(function_response_events) > 1:
     # 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)
-    with tracer.start_as_current_span('tool_response'):
-      trace_tool_response(
-          invocation_context=invocation_context,
-          event_id=merged_event.id,
+    with tracer.start_as_current_span('execute_tool (merged)'):
+      trace_merged_tool_calls(
+          response_event_id=merged_event.id,
           function_response_event=merged_event,
       )
   return merged_event
@@ -225,65 +232,81 @@ async def handle_function_calls_live(
     tool, tool_context = _get_tool_and_context(
         invocation_context, function_call_event, function_call, tools_dict
     )
-    # do not use "args" as the variable name, because it is a reserved keyword
-    # in python debugger.
-    function_args = function_call.args or {}
-    function_response = None
-    # # Calls the tool if before_tool_callback does not exist or returns None.
-    # if agent.before_tool_callback:
-    #   function_response = agent.before_tool_callback(
-    #       tool, function_args, tool_context
-    #   )
-    if agent.before_tool_callback:
-      function_response = agent.before_tool_callback(
-          tool=tool, args=function_args, tool_context=tool_context
-      )
-      if inspect.isawaitable(function_response):
-        function_response = await function_response
+    with tracer.start_as_current_span(f'execute_tool {tool.name}'):
+      # do not use "args" as the variable name, because it is a reserved keyword
+      # in python debugger.
+      function_args = function_call.args or {}
+      function_response = None
+      # # Calls the tool if before_tool_callback does not exist or returns None.
+      # if agent.before_tool_callback:
+      #   function_response = agent.before_tool_callback(
+      #       tool, function_args, tool_context
+      #   )
+      if agent.before_tool_callback:
+        function_response = agent.before_tool_callback(
+            tool=tool, args=function_args, tool_context=tool_context
+        )
+        if inspect.isawaitable(function_response):
+          function_response = await function_response
 
-    if not function_response:
-      function_response = await _process_function_live_helper(
-          tool, tool_context, function_call, function_args, invocation_context
-      )
+      if not function_response:
+        function_response = await _process_function_live_helper(
+            tool, tool_context, function_call, function_args, invocation_context
+        )
 
-    # Calls after_tool_callback if it exists.
-    # if agent.after_tool_callback:
-    #   new_response = agent.after_tool_callback(
-    #       tool,
-    #       function_args,
-    #       tool_context,
-    #       function_response,
-    #   )
-    #   if new_response:
-    #     function_response = new_response
-    if agent.after_tool_callback:
-      altered_function_response = agent.after_tool_callback(
+      # Calls after_tool_callback if it exists.
+      # if agent.after_tool_callback:
+      #   new_response = agent.after_tool_callback(
+      #       tool,
+      #       function_args,
+      #       tool_context,
+      #       function_response,
+      #   )
+      #   if new_response:
+      #     function_response = new_response
+      if agent.after_tool_callback:
+        altered_function_response = agent.after_tool_callback(
+            tool=tool,
+            args=function_args,
+            tool_context=tool_context,
+            tool_response=function_response,
+        )
+        if inspect.isawaitable(altered_function_response):
+          altered_function_response = await altered_function_response
+        if altered_function_response is not None:
+          function_response = altered_function_response
+
+      if tool.is_long_running:
+        # Allow async function to return None to not provide function response.
+        if not function_response:
+          continue
+
+      # Builds the function response event.
+      function_response_event = __build_response_event(
+          tool, function_response, tool_context, invocation_context
+      )
+      trace_tool_call(
           tool=tool,
           args=function_args,
-          tool_context=tool_context,
-          tool_response=function_response,
+          response_event_id=function_response_event.id,
+          function_response=function_response,
       )
-      if inspect.isawaitable(altered_function_response):
-        altered_function_response = await altered_function_response
-      if altered_function_response is not None:
-        function_response = altered_function_response
-
-    if tool.is_long_running:
-      # Allow async function to return None to not provide function response.
-      if not function_response:
-        continue
-
-    # Builds the function response event.
-    function_response_event = __build_response_event(
-        tool, function_response, tool_context, invocation_context
-    )
-    function_response_events.append(function_response_event)
+      function_response_events.append(function_response_event)
 
   if not function_response_events:
     return None
   merged_event = merge_parallel_function_response_events(
       function_response_events
   )
+  if len(function_response_events) > 1:
+    # 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)
+    with tracer.start_as_current_span('execute_tool (merged)'):
+      trace_merged_tool_calls(
+          response_event_id=merged_event.id,
+          function_response_event=merged_event,
+      )
   return merged_event
 
 
@@ -410,14 +433,12 @@ async def __call_tool_live(
     invocation_context: InvocationContext,
 ) -> AsyncGenerator[Event, None]:
   """Calls the tool asynchronously (awaiting the coroutine)."""
-  with tracer.start_as_current_span(f'tool_call [{tool.name}]'):
-    trace_tool_call(args=args)
-    async for item in tool._call_live(
-        args=args,
-        tool_context=tool_context,
-        invocation_context=invocation_context,
-    ):
-      yield item
+  async for item in tool._call_live(
+      args=args,
+      tool_context=tool_context,
+      invocation_context=invocation_context,
+  ):
+    yield item
 
 
 async def __call_tool_async(
@@ -426,9 +447,7 @@ async def __call_tool_async(
     tool_context: ToolContext,
 ) -> Any:
   """Calls the tool."""
-  with tracer.start_as_current_span(f'tool_call [{tool.name}]'):
-    trace_tool_call(args=args)
-    return await tool.run_async(args=args, tool_context=tool_context)
+  return await tool.run_async(args=args, tool_context=tool_context)
 
 
 def __build_response_event(
@@ -437,35 +456,29 @@ def __build_response_event(
     tool_context: ToolContext,
     invocation_context: InvocationContext,
 ) -> Event:
-  with tracer.start_as_current_span(f'tool_response [{tool.name}]'):
-    # Specs requires the result to be a dict.
-    if not isinstance(function_result, dict):
-      function_result = {'result': function_result}
+  # Specs requires the result to be a dict.
+  if not isinstance(function_result, dict):
+    function_result = {'result': function_result}
 
-    part_function_response = types.Part.from_function_response(
-        name=tool.name, response=function_result
-    )
-    part_function_response.function_response.id = tool_context.function_call_id
+  part_function_response = types.Part.from_function_response(
+      name=tool.name, response=function_result
+  )
+  part_function_response.function_response.id = tool_context.function_call_id
 
-    content = types.Content(
-        role='user',
-        parts=[part_function_response],
-    )
+  content = types.Content(
+      role='user',
+      parts=[part_function_response],
+  )
 
-    function_response_event = Event(
-        invocation_id=invocation_context.invocation_id,
-        author=invocation_context.agent.name,
-        content=content,
-        actions=tool_context.actions,
-        branch=invocation_context.branch,
-    )
+  function_response_event = Event(
+      invocation_id=invocation_context.invocation_id,
+      author=invocation_context.agent.name,
+      content=content,
+      actions=tool_context.actions,
+      branch=invocation_context.branch,
+  )
 
-    trace_tool_response(
-        invocation_context=invocation_context,
-        event_id=function_response_event.id,
-        function_response_event=function_response_event,
-    )
-    return function_response_event
+  return function_response_event
 
 
 def merge_parallel_function_response_events(

google/adk/memory/vertex_ai_rag_memory_service.py

@@ -124,8 +124,8 @@ class VertexAiRagMemoryService(BaseMemoryService):
     for context in response.contexts.contexts:
       # filter out context that is not related
       # TODO: Add server side filtering by app_name and user_id.
-      # if not context.source_display_name.startswith(f"{app_name}.{user_id}."):
-      #   continue
+      if not context.source_display_name.startswith(f"{app_name}.{user_id}."):
+        continue
       session_id = context.source_display_name.split(".")[-1]
       events = []
       if context.text:

google/adk/models/anthropic_llm.py

@@ -135,6 +135,10 @@ def content_block_to_part(
 def message_to_generate_content_response(
     message: anthropic_types.Message,
 ) -> LlmResponse:
+  logger.info(
+      "Claude response: %s",
+      message.model_dump_json(indent=2, exclude_none=True),
+  )
 
   return LlmResponse(
       content=types.Content(
@@ -208,7 +212,7 @@ class Claude(BaseLlm):
   @staticmethod
   @override
   def supported_models() -> list[str]:
-    return [r"claude-3-.*"]
+    return [r"claude-3-.*", r"claude-.*-4.*"]
 
   @override
   async def generate_content_async(
@@ -229,14 +233,11 @@ class Claude(BaseLlm):
           for tool in llm_request.config.tools[0].function_declarations
       ]
     tool_choice = (
-        anthropic_types.ToolChoiceAutoParam(
-            type="auto",
-            # TODO: allow parallel tool use.
-            disable_parallel_tool_use=True,
-        )
+        anthropic_types.ToolChoiceAutoParam(type="auto")
         if llm_request.tools_dict
         else NOT_GIVEN
     )
+    # TODO(b/421255973): Enable streaming for anthropic models.
     message = self._anthropic_client.messages.create(
         model=llm_request.model,
         system=llm_request.config.system_instruction,
@@ -245,10 +246,6 @@ class Claude(BaseLlm):
         tool_choice=tool_choice,
         max_tokens=MAX_TOKEN,
     )
-    logger.info(
-        "Claude response: %s",
-        message.model_dump_json(indent=2, exclude_none=True),
-    )
     yield message_to_generate_content_response(message)
 
   @cached_property

google/adk/models/google_llm.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 import contextlib
 from functools import cached_property
 import logging
+import os
 import sys
 from typing import AsyncGenerator
 from typing import cast
@@ -28,6 +29,7 @@ from google.genai import types
 from typing_extensions import override
 
 from .. import version
+from ..utils.variant_utils import GoogleLLMVariant
 from .base_llm import BaseLlm
 from .base_llm_connection import BaseLlmConnection
 from .gemini_llm_connection import GeminiLlmConnection
@@ -40,6 +42,8 @@ logger = logging.getLogger('google_adk.' + __name__)
 
 _NEW_LINE = '\n'
 _EXCLUDED_PART_FIELD = {'inline_data': {'data'}}
+_AGENT_ENGINE_TELEMETRY_TAG = 'remote_reasoning_engine'
+_AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_AGENT_ENGINE_ID'
 
 
 class Gemini(BaseLlm):
@@ -80,7 +84,7 @@ class Gemini(BaseLlm):
     Yields:
       LlmResponse: The model response.
     """
-
+    self._preprocess_request(llm_request)
     self._maybe_append_user_content(llm_request)
     logger.info(
         'Sending out request, model: %s, backend: %s, stream: %s',
@@ -97,6 +101,7 @@ class Gemini(BaseLlm):
           config=llm_request.config,
       )
       response = None
+      thought_text = ''
       text = ''
       usage_metadata = None
       # for sse, similar as bidi (see receive method in gemini_llm_connecton.py),
@@ -113,32 +118,43 @@ class Gemini(BaseLlm):
             and llm_response.content.parts
             and llm_response.content.parts[0].text
         ):
-          text += llm_response.content.parts[0].text
+          part0 = llm_response.content.parts[0]
+          if part0.thought:
+            thought_text += part0.text
+          else:
+            text += part0.text
           llm_response.partial = True
-        elif text and (
+        elif (thought_text or text) and (
             not llm_response.content
             or not llm_response.content.parts
             # don't yield the merged text event when receiving audio data
             or not llm_response.content.parts[0].inline_data
         ):
+          parts = []
+          if thought_text:
+            parts.append(types.Part(text=thought_text, thought=True))
+          if text:
+            parts.append(types.Part.from_text(text=text))
           yield LlmResponse(
-              content=types.ModelContent(
-                  parts=[types.Part.from_text(text=text)],
-              ),
-              usage_metadata=usage_metadata,
+              content=types.ModelContent(parts=parts),
+              usage_metadata=llm_response.usage_metadata,
           )
+          thought_text = ''
           text = ''
         yield llm_response
       if (
-          text
+          (text or thought_text)
           and response
           and response.candidates
           and response.candidates[0].finish_reason == types.FinishReason.STOP
       ):
+        parts = []
+        if thought_text:
+          parts.append(types.Part(text=thought_text, thought=True))
+        if text:
+          parts.append(types.Part.from_text(text=text))
         yield LlmResponse(
-            content=types.ModelContent(
-                parts=[types.Part.from_text(text=text)],
-            ),
+            content=types.ModelContent(parts=parts),
             usage_metadata=usage_metadata,
         )
 
@@ -163,12 +179,18 @@ class Gemini(BaseLlm):
     )
 
   @cached_property
-  def _api_backend(self) -> str:
-    return 'vertex' if self.api_client.vertexai else 'ml_dev'
+  def _api_backend(self) -> GoogleLLMVariant:
+    return (
+        GoogleLLMVariant.VERTEX_AI
+        if self.api_client.vertexai
+        else GoogleLLMVariant.GEMINI_API
+    )
 
   @cached_property
   def _tracking_headers(self) -> dict[str, str]:
     framework_label = f'google-adk/{version.__version__}'
+    if os.environ.get(_AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME):
+      framework_label = f'{framework_label}+{_AGENT_ENGINE_TELEMETRY_TAG}'
     language_label = 'gl-python/' + sys.version.split()[0]
     version_header_value = f'{framework_label} {language_label}'
     tracking_headers = {
@@ -179,7 +201,7 @@ class Gemini(BaseLlm):
 
   @cached_property
   def _live_api_client(self) -> Client:
-    if self._api_backend == 'vertex':
+    if self._api_backend == GoogleLLMVariant.VERTEX_AI:
       # use beta version for vertex api
       api_version = 'v1beta1'
       # use default api version for vertex
@@ -189,7 +211,7 @@ class Gemini(BaseLlm):
           )
       )
     else:
-      # use v1alpha for ml_dev
+      # use v1alpha for using API KEY from Google AI Studio
       api_version = 'v1alpha'
       return Client(
           http_options=types.HttpOptions(
@@ -220,6 +242,12 @@ class Gemini(BaseLlm):
     ) as live_session:
       yield GeminiLlmConnection(live_session)
 
+  def _preprocess_request(self, llm_request: LlmRequest) -> None:
+
+    if llm_request.config and self._api_backend == GoogleLLMVariant.GEMINI_API:
+      # Using API key from Google AI Studio to call model doesn't support labels.
+      llm_request.config.labels = None
+
 
 def _build_function_declaration_log(
     func_decl: types.FunctionDeclaration,
@@ -230,10 +258,10 @@ def _build_function_declaration_log(
         k: v.model_dump(exclude_none=True)
         for k, v in func_decl.parameters.properties.items()
     })
-  return_str = 'None'
+  return_str = ''
   if func_decl.response:
-    return_str = str(func_decl.response.model_dump(exclude_none=True))
-  return f'{func_decl.name}: {param_str} -> {return_str}'
+    return_str = '-> ' + str(func_decl.response.model_dump(exclude_none=True))
+  return f'{func_decl.name}: {param_str} {return_str}'
 
 
 def _build_request_log(req: LlmRequest) -> str: