google-adk 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

google/adk/evaluation/agent_evaluator.py

@@ -76,7 +76,7 @@ class AgentEvaluator:
     return DEFAULT_CRITERIA
 
   @staticmethod
-  def evaluate(
+  async def evaluate(
       agent_module,
       eval_dataset_file_path_or_dir,
       num_runs=NUM_RUNS,
@@ -120,7 +120,7 @@ class AgentEvaluator:
 
       AgentEvaluator._validate_input([dataset], criteria)
 
-      evaluation_response = AgentEvaluator._generate_responses(
+      evaluation_response = await AgentEvaluator._generate_responses(
           agent_module,
           [dataset],
           num_runs,
@@ -246,7 +246,7 @@ class AgentEvaluator:
     return inferred_criteria
 
   @staticmethod
-  def _generate_responses(
+  async def _generate_responses(
       agent_module, eval_dataset, num_runs, agent_name=None, initial_session={}
   ):
     """Generates evaluation responses by running the agent module multiple times."""

google/adk/evaluation/evaluation_constants.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 class EvalConstants:
   """Holds constants for evaluation file constants."""
 

google/adk/evaluation/evaluation_generator.py

@@ -32,7 +32,7 @@ class EvaluationGenerator:
   """Generates evaluation responses for agents."""
 
   @staticmethod
-  def generate_responses(
+  async def generate_responses(
       eval_dataset,
       agent_module_path,
       repeat_num=3,
@@ -107,7 +107,7 @@ class EvaluationGenerator:
     )
 
   @staticmethod
-  def _process_query_with_root_agent(
+  async def _process_query_with_root_agent(
       data,
       root_agent,
       reset_func,
@@ -128,7 +128,7 @@ class EvaluationGenerator:
           all_mock_tools.add(expected[EvalConstants.TOOL_NAME])
 
     eval_data_copy = data.copy()
-    EvaluationGenerator.apply_before_tool_callback(
+    await EvaluationGenerator.apply_before_tool_callback(
         root_agent,
         lambda *args: EvaluationGenerator.before_tool_callback(
             *args, eval_dataset=eval_data_copy
@@ -247,7 +247,7 @@ class EvaluationGenerator:
     return None
 
   @staticmethod
-  def apply_before_tool_callback(
+  async def apply_before_tool_callback(
       agent: BaseAgent,
       callback: BeforeToolCallback,
       all_mock_tools: set[str],
@@ -265,6 +265,6 @@ class EvaluationGenerator:
 
     # Apply recursively to subagents if they exist
     for sub_agent in agent.sub_agents:
-      EvaluationGenerator.apply_before_tool_callback(
+      await EvaluationGenerator.apply_before_tool_callback(
           sub_agent, callback, all_mock_tools
       )

google/adk/evaluation/response_evaluator.py

@@ -28,7 +28,7 @@ class ResponseEvaluator:
       raw_eval_dataset: list[list[dict[str, Any]]],
       evaluation_criteria: list[str],
       *,
-      print_detailed_results: bool = False
+      print_detailed_results: bool = False,
   ):
     r"""Returns the value of requested evaluation metrics.
 

google/adk/events/event.py

@@ -48,6 +48,7 @@ class Event(LlmResponse):
   model_config = ConfigDict(
       extra='forbid', ser_json_bytes='base64', val_json_bytes='base64'
   )
+  """The pydantic model config."""
 
   # TODO: revert to be required after spark migration
   invocation_id: str = ''

google/adk/events/event_actions.py

@@ -27,6 +27,7 @@ class EventActions(BaseModel):
   """Represents the actions attached to an event."""
 
   model_config = ConfigDict(extra='forbid')
+  """The pydantic model config."""
 
   skip_summarization: Optional[bool] = None
   """If true, it won't call model to summarize function response.
@@ -48,8 +49,13 @@ class EventActions(BaseModel):
   """The agent is escalating to a higher level agent."""
 
   requested_auth_configs: dict[str, AuthConfig] = Field(default_factory=dict)
-  """Will only be set by a tool response indicating tool request euc.
-  dict key is the function call id since one function call response (from model)
-  could correspond to multiple function calls.
-  dict value is the required auth config.
+  """Authentication configurations requested by tool responses.
+
+  This field will only be set by a tool response event indicating tool request
+  auth credential.
+  - Keys: The function call id. Since one function response event could contain
+  multiple function responses that correspond to multiple function calls. Each
+  function call could request different auth configs. This id is used to
+  identify the function call.
+  - Values: The requested auth config.
   """

google/adk/examples/example.py

@@ -23,5 +23,6 @@ class Example(BaseModel):
     input: The input content for the example.
     output: The expected output content for the example.
   """
+
   input: types.Content
   output: list[types.Content]

google/adk/flows/__init__.py

@@ -11,4 +11,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-

google/adk/flows/llm_flows/_code_execution.py

@@ -122,7 +122,7 @@ class _CodeExecutionRequestProcessor(BaseLlmRequestProcessor):
     if not invocation_context.agent.code_executor:
       return
 
-    for event in _run_pre_processor(invocation_context, llm_request):
+    async for event in _run_pre_processor(invocation_context, llm_request):
       yield event
 
     # Convert the code execution parts to text parts.
@@ -152,17 +152,17 @@ class _CodeExecutionResponseProcessor(BaseLlmResponseProcessor):
     if llm_response.partial:
       return
 
-    for event in _run_post_processor(invocation_context, llm_response):
+    async for event in _run_post_processor(invocation_context, llm_response):
       yield event
 
 
 response_processor = _CodeExecutionResponseProcessor()
 
 
-def _run_pre_processor(
+async def _run_pre_processor(
     invocation_context: InvocationContext,
     llm_request: LlmRequest,
-) -> Generator[Event, None, None]:
+) -> AsyncGenerator[Event, None]:
   """Pre-process the user message by adding the user message to the Colab notebook."""
   from ...agents.llm_agent import LlmAgent
 
@@ -242,17 +242,17 @@ def _run_pre_processor(
     code_executor_context.add_processed_file_names([file.name])
 
     # Emit the execution result, and add it to the LLM request.
-    execution_result_event = _post_process_code_execution_result(
+    execution_result_event = await _post_process_code_execution_result(
         invocation_context, code_executor_context, code_execution_result
     )
     yield execution_result_event
     llm_request.contents.append(copy.deepcopy(execution_result_event.content))
 
 
-def _run_post_processor(
+async def _run_post_processor(
     invocation_context: InvocationContext,
     llm_response,
-) -> Generator[Event, None, None]:
+) -> AsyncGenerator[Event, None]:
   """Post-process the model response by extracting and executing the first code block."""
   agent = invocation_context.agent
   code_executor = agent.code_executor
@@ -305,7 +305,7 @@ def _run_post_processor(
       code_execution_result.stdout,
       code_execution_result.stderr,
   )
-  yield _post_process_code_execution_result(
+  yield await _post_process_code_execution_result(
       invocation_context, code_executor_context, code_execution_result
   )
 
@@ -375,7 +375,7 @@ def _get_or_set_execution_id(
   return execution_id
 
 
-def _post_process_code_execution_result(
+async def _post_process_code_execution_result(
     invocation_context: InvocationContext,
     code_executor_context: CodeExecutorContext,
     code_execution_result: CodeExecutionResult,
@@ -406,7 +406,7 @@ def _post_process_code_execution_result(
 
   # Handle output files.
   for output_file in code_execution_result.output_files:
-    version = invocation_context.artifact_service.save_artifact(
+    version = await invocation_context.artifact_service.save_artifact(
         app_name=invocation_context.app_name,
         user_id=invocation_context.user_id,
         session_id=invocation_context.session.id,

google/adk/flows/llm_flows/base_llm_flow.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 from abc import ABC
 import asyncio
+import inspect
 import logging
 from typing import AsyncGenerator
 from typing import cast
@@ -190,6 +191,17 @@ class BaseLlmFlow(ABC):
       llm_request: LlmRequest,
   ) -> AsyncGenerator[Event, None]:
     """Receive data from model and process events using BaseLlmConnection."""
+    def get_author(llm_response):
+      """Get the author of the event.
+
+      When the model returns transcription, the author is "user". Otherwise, the
+      author is the agent.
+      """
+      if llm_response and llm_response.content and llm_response.content.role == "user":
+        return "user"
+      else:
+        return invocation_context.agent.name
+
     assert invocation_context.live_request_queue
     try:
       while True:
@@ -197,7 +209,7 @@ class BaseLlmFlow(ABC):
           model_response_event = Event(
               id=Event.new_id(),
               invocation_id=invocation_context.invocation_id,
-              author=invocation_context.agent.name,
+              author=get_author(llm_response),
           )
           async for event in self._postprocess_live(
               invocation_context,
@@ -249,7 +261,6 @@ class BaseLlmFlow(ABC):
 
     # Calls the LLM.
     model_response_event = Event(
-        id=Event.new_id(),
         invocation_id=invocation_context.invocation_id,
         author=invocation_context.agent.name,
         branch=invocation_context.branch,
@@ -261,6 +272,8 @@ class BaseLlmFlow(ABC):
       async for event in self._postprocess_async(
           invocation_context, llm_request, llm_response, model_response_event
       ):
+        # Use a new id for every event.
+        event.id = Event.new_id()
         yield event
 
   async def _preprocess_async(
@@ -437,7 +450,7 @@ class BaseLlmFlow(ABC):
       model_response_event: Event,
   ) -> AsyncGenerator[LlmResponse, None]:
     # Runs before_model_callback if it exists.
-    if response := self._handle_before_model_callback(
+    if response := await self._handle_before_model_callback(
         invocation_context, llm_request, model_response_event
     ):
       yield response
@@ -450,7 +463,7 @@ class BaseLlmFlow(ABC):
         invocation_context.live_request_queue = LiveRequestQueue()
         async for llm_response in self.run_live(invocation_context):
           # Runs after_model_callback if it exists.
-          if altered_llm_response := self._handle_after_model_callback(
+          if altered_llm_response := await self._handle_after_model_callback(
               invocation_context, llm_response, model_response_event
           ):
             llm_response = altered_llm_response
@@ -479,14 +492,14 @@ class BaseLlmFlow(ABC):
               llm_response,
           )
           # Runs after_model_callback if it exists.
-          if altered_llm_response := self._handle_after_model_callback(
+          if altered_llm_response := await self._handle_after_model_callback(
               invocation_context, llm_response, model_response_event
           ):
             llm_response = altered_llm_response
 
           yield llm_response
 
-  def _handle_before_model_callback(
+  async def _handle_before_model_callback(
       self,
       invocation_context: InvocationContext,
       llm_request: LlmRequest,
@@ -498,17 +511,23 @@ class BaseLlmFlow(ABC):
     if not isinstance(agent, LlmAgent):
       return
 
-    if not agent.before_model_callback:
+    if not agent.canonical_before_model_callbacks:
       return
 
     callback_context = CallbackContext(
         invocation_context, event_actions=model_response_event.actions
     )
-    return agent.before_model_callback(
-        callback_context=callback_context, llm_request=llm_request
-    )
 
-  def _handle_after_model_callback(
+    for callback in agent.canonical_before_model_callbacks:
+      before_model_callback_content = callback(
+          callback_context=callback_context, llm_request=llm_request
+      )
+      if inspect.isawaitable(before_model_callback_content):
+        before_model_callback_content = await before_model_callback_content
+      if before_model_callback_content:
+        return before_model_callback_content
+
+  async def _handle_after_model_callback(
       self,
       invocation_context: InvocationContext,
       llm_response: LlmResponse,
@@ -520,15 +539,21 @@ class BaseLlmFlow(ABC):
     if not isinstance(agent, LlmAgent):
       return
 
-    if not agent.after_model_callback:
+    if not agent.canonical_after_model_callbacks:
       return
 
     callback_context = CallbackContext(
         invocation_context, event_actions=model_response_event.actions
     )
-    return agent.after_model_callback(
-        callback_context=callback_context, llm_response=llm_response
-    )
+
+    for callback in agent.canonical_after_model_callbacks:
+      after_model_callback_content = callback(
+          callback_context=callback_context, llm_response=llm_response
+      )
+      if inspect.isawaitable(after_model_callback_content):
+        after_model_callback_content = await after_model_callback_content
+      if after_model_callback_content:
+        return after_model_callback_content
 
   def _finalize_model_response_event(
       self,

google/adk/flows/llm_flows/basic.py

@@ -62,6 +62,9 @@ class _BasicLlmRequestProcessor(BaseLlmRequestProcessor):
     llm_request.live_connect_config.output_audio_transcription = (
         invocation_context.run_config.output_audio_transcription
     )
+    llm_request.live_connect_config.input_audio_transcription = (
+        invocation_context.run_config.input_audio_transcription
+    )
 
     # TODO: handle tool append here, instead of in BaseTool.process_llm_request.
 

google/adk/flows/llm_flows/contents.py

@@ -15,9 +15,7 @@
 from __future__ import annotations
 
 import copy
-from typing import AsyncGenerator
-from typing import Generator
-from typing import Optional
+from typing import AsyncGenerator, Generator, Optional
 
 from google.genai import types
 from typing_extensions import override
@@ -202,8 +200,14 @@ def _get_contents(
   # Parse the events, leaving the contents and the function calls and
   # responses from the current agent.
   for event in events:
-    if not event.content or not event.content.role:
-      # Skip events without content, or generated neither by user nor by model.
+    if (
+        not event.content
+        or not event.content.role
+        or not event.content.parts
+        or event.content.parts[0].text == ''
+    ):
+      # Skip events without content, or generated neither by user nor by model
+      # or has empty text.
       # E.g. events purely for mutating session states.
       continue
     if not _is_event_belongs_to_branch(current_branch, event):

google/adk/flows/llm_flows/functions.py

@@ -151,28 +151,33 @@ async def handle_function_calls_async(
     # 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.
+    function_response: Optional[dict] = None
+
+    # before_tool_callback (sync or async)
     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 __call_tool_async(
           tool, args=function_args, tool_context=tool_context
       )
 
-    # Calls after_tool_callback if it exists.
+    # after_tool_callback (sync or async)
     if agent.after_tool_callback:
-      new_response = 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 new_response:
-        function_response = new_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 long running function to return None to not provide function response.
@@ -223,11 +228,17 @@ async def handle_function_calls_live(
     # 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.
+    # # 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, function_args, tool_context
+          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(
@@ -235,15 +246,26 @@ async def handle_function_calls_live(
       )
 
     # 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:
-      new_response = agent.after_tool_callback(
-          tool,
-          function_args,
-          tool_context,
-          function_response,
+      altered_function_response = agent.after_tool_callback(
+          tool=tool,
+          args=function_args,
+          tool_context=tool_context,
+          tool_response=function_response,
       )
-      if new_response:
-        function_response = new_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.
@@ -310,7 +332,7 @@ async def _process_function_live_helper(
       function_response = {
           'status': f'No active streaming function named {function_name} found'
       }
-  elif hasattr(tool, "func") and inspect.isasyncgenfunction(tool.func):
+  elif hasattr(tool, 'func') and inspect.isasyncgenfunction(tool.func):
     # for streaming tool use case
     # we require the function to be a async generator function
     async def run_tool_and_update_queue(tool, function_args, tool_context):

google/adk/flows/llm_flows/instructions.py

@@ -56,13 +56,13 @@ class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
       raw_si = root_agent.canonical_global_instruction(
           ReadonlyContext(invocation_context)
       )
-      si = _populate_values(raw_si, invocation_context)
+      si = await _populate_values(raw_si, invocation_context)
       llm_request.append_instructions([si])
 
     # Appends agent instructions if set.
     if agent.instruction:  # not empty str
       raw_si = agent.canonical_instruction(ReadonlyContext(invocation_context))
-      si = _populate_values(raw_si, invocation_context)
+      si = await _populate_values(raw_si, invocation_context)
       llm_request.append_instructions([si])
 
     # Maintain async generator behavior
@@ -73,13 +73,24 @@ class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
 request_processor = _InstructionsLlmRequestProcessor()
 
 
-def _populate_values(
+async def _populate_values(
     instruction_template: str,
     context: InvocationContext,
 ) -> str:
   """Populates values in the instruction template, e.g. state, artifact, etc."""
 
-  def _replace_match(match) -> str:
+  async def _async_sub(pattern, repl_async_fn, string) -> str:
+    result = []
+    last_end = 0
+    for match in re.finditer(pattern, string):
+      result.append(string[last_end : match.start()])
+      replacement = await repl_async_fn(match)
+      result.append(replacement)
+      last_end = match.end()
+    result.append(string[last_end:])
+    return ''.join(result)
+
+  async def _replace_match(match) -> str:
     var_name = match.group().lstrip('{').rstrip('}').strip()
     optional = False
     if var_name.endswith('?'):
@@ -89,7 +100,7 @@ def _populate_values(
       var_name = var_name.removeprefix('artifact.')
       if context.artifact_service is None:
         raise ValueError('Artifact service is not initialized.')
-      artifact = context.artifact_service.load_artifact(
+      artifact = await context.artifact_service.load_artifact(
           app_name=context.session.app_name,
           user_id=context.session.user_id,
           session_id=context.session.id,
@@ -109,7 +120,7 @@ def _populate_values(
         else:
           raise KeyError(f'Context variable not found: `{var_name}`.')
 
-  return re.sub(r'{+[^{}]*}+', _replace_match, instruction_template)
+  return await _async_sub(r'{+[^{}]*}+', _replace_match, instruction_template)
 
 
 def _is_valid_state_name(var_name):

google/adk/memory/base_memory_service.py

@@ -28,6 +28,7 @@ class MemoryResult(BaseModel):
       session_id: The session id associated with the memory.
       events: A list of events in the session.
   """
+
   session_id: str
   events: list[Event]
 
@@ -38,6 +39,7 @@ class SearchMemoryResponse(BaseModel):
   Attributes:
       memories: A list of memory results matching the search query.
   """
+
   memories: list[MemoryResult] = Field(default_factory=list)
 
 
@@ -49,7 +51,7 @@ class BaseMemoryService(abc.ABC):
   """
 
   @abc.abstractmethod
-  def add_session_to_memory(self, session: Session):
+  async def add_session_to_memory(self, session: Session):
     """Adds a session to the memory service.
 
     A session may be added multiple times during its lifetime.
@@ -59,7 +61,7 @@ class BaseMemoryService(abc.ABC):
     """
 
   @abc.abstractmethod
-  def search_memory(
+  async def search_memory(
       self, *, app_name: str, user_id: str, query: str
   ) -> SearchMemoryResponse:
     """Searches for sessions that match the query.

google/adk/memory/in_memory_memory_service.py

@@ -29,13 +29,13 @@ class InMemoryMemoryService(BaseMemoryService):
     self.session_events: dict[str, list[Event]] = {}
     """keys are app_name/user_id/session_id"""
 
-  def add_session_to_memory(self, session: Session):
+  async def add_session_to_memory(self, session: Session):
     key = f'{session.app_name}/{session.user_id}/{session.id}'
     self.session_events[key] = [
         event for event in session.events if event.content
     ]
 
-  def search_memory(
+  async def search_memory(
       self, *, app_name: str, user_id: str, query: str
   ) -> SearchMemoryResponse:
     """Prototyping purpose only."""

google/adk/memory/vertex_ai_rag_memory_service.py

@@ -54,7 +54,7 @@ class VertexAiRagMemoryService(BaseMemoryService):
     )
 
   @override
-  def add_session_to_memory(self, session: Session):
+  async def add_session_to_memory(self, session: Session):
     with tempfile.NamedTemporaryFile(
         mode="w", delete=False, suffix=".txt"
     ) as temp_file:
@@ -91,7 +91,7 @@ class VertexAiRagMemoryService(BaseMemoryService):
     os.remove(temp_file_path)
 
   @override
-  def search_memory(
+  async def search_memory(
       self, *, app_name: str, user_id: str, query: str
   ) -> SearchMemoryResponse:
     """Searches for sessions that match the query using rag.retrieval_query."""

google/adk/models/anthropic_llm.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 from functools import cached_property
 import logging
 import os
+from typing import Any
 from typing import AsyncGenerator
 from typing import Generator
 from typing import Iterable
@@ -151,6 +152,24 @@ def message_to_generate_content_response(
   )
 
 
+def _update_type_string(value_dict: dict[str, Any]):
+  """Updates 'type' field to expected JSON schema format."""
+  if "type" in value_dict:
+    value_dict["type"] = value_dict["type"].lower()
+
+  if "items" in value_dict:
+    # 'type' field could exist for items as well, this would be the case if
+    # items represent primitive types.
+    _update_type_string(value_dict["items"])
+
+    if "properties" in value_dict["items"]:
+      # There could be properties as well on the items, especially if the items
+      # are complex object themselves. We recursively traverse each individual
+      # property as well and fix the "type" value.
+      for _, value in value_dict["items"]["properties"].items():
+        _update_type_string(value)
+
+
 def function_declaration_to_tool_param(
     function_declaration: types.FunctionDeclaration,
 ) -> anthropic_types.ToolParam:
@@ -163,8 +182,7 @@ def function_declaration_to_tool_param(
   ):
     for key, value in function_declaration.parameters.properties.items():
       value_dict = value.model_dump(exclude_none=True)
-      if "type" in value_dict:
-        value_dict["type"] = value_dict["type"].lower()
+      _update_type_string(value_dict)
       properties[key] = value_dict
 
   return anthropic_types.ToolParam(