Merge branch 'main' into raflFaisal/generaliseContainerEnvForAgent

Author: raflFaisal

.github/workflows/python-unit-tests.yml

@@ -36,6 +36,5 @@ jobs:
         source .venv/bin/activate
         pytest tests/unittests \
           --ignore=tests/unittests/artifacts/test_artifact_service.py \
-          --ignore=tests/unittests/tools/application_integration_tool/clients/test_connections_client.py \
           --ignore=tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py        
           

src/google/adk/agents/base_agent.py

@@ -14,7 +14,8 @@
 
 from __future__ import annotations
 
-from typing import Any
+import inspect
+from typing import Any, Awaitable, Union
 from typing import AsyncGenerator
 from typing import Callable
 from typing import final
@@ -37,10 +38,15 @@
 
 tracer = trace.get_tracer('gcp.vertex.agent')
 
-BeforeAgentCallback = Callable[[CallbackContext], Optional[types.Content]]
+BeforeAgentCallback = Callable[
+    [CallbackContext],
+    Union[Awaitable[Optional[types.Content]], Optional[types.Content]],
+]
 
-
-AfterAgentCallback = Callable[[CallbackContext], Optional[types.Content]]
+AfterAgentCallback = Callable[
+    [CallbackContext],
+    Union[Awaitable[Optional[types.Content]], Optional[types.Content]],
+]
 
 
 class BaseAgent(BaseModel):
@@ -119,7 +125,7 @@ async def run_async(
     with tracer.start_as_current_span(f'agent_run [{self.name}]'):
       ctx = self._create_invocation_context(parent_context)
 
-      if event := self.__handle_before_agent_callback(ctx):
+      if event := await self.__handle_before_agent_callback(ctx):
         yield event
       if ctx.end_invocation:
         return
@@ -130,7 +136,7 @@ async def run_async(
       if ctx.end_invocation:
         return
 
-      if event := self.__handle_after_agent_callback(ctx):
+      if event := await self.__handle_after_agent_callback(ctx):
         yield event
 
   @final
@@ -230,7 +236,7 @@ def _create_invocation_context(
       invocation_context.branch = f'{parent_context.branch}.{self.name}'
     return invocation_context
 
-  def __handle_before_agent_callback(
+  async def __handle_before_agent_callback(
       self, ctx: InvocationContext
   ) -> Optional[Event]:
     """Runs the before_agent_callback if it exists.
@@ -248,6 +254,9 @@ def __handle_before_agent_callback(
         callback_context=callback_context
     )
 
+    if inspect.isawaitable(before_agent_callback_content):
+      before_agent_callback_content = await before_agent_callback_content
+
     if before_agent_callback_content:
       ret_event = Event(
           invocation_id=ctx.invocation_id,
@@ -269,7 +278,7 @@ def __handle_before_agent_callback(
 
     return ret_event
 
-  def __handle_after_agent_callback(
+  async def __handle_after_agent_callback(
       self, invocation_context: InvocationContext
   ) -> Optional[Event]:
     """Runs the after_agent_callback if it exists.
@@ -287,6 +296,9 @@ def __handle_after_agent_callback(
         callback_context=callback_context
     )
 
+    if inspect.isawaitable(after_agent_callback_content):
+      after_agent_callback_content = await after_agent_callback_content
+
     if after_agent_callback_content or callback_context.state.has_delta():
       ret_event = Event(
           invocation_id=invocation_context.invocation_id,

src/google/adk/agents/llm_agent.py

@@ -47,14 +47,26 @@
 
 logger = logging.getLogger(__name__)
 
+_SingleBeforeModelCallback: TypeAlias = Callable[
+    [CallbackContext, LlmRequest],
+    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
+]
 
-BeforeModelCallback: TypeAlias = Callable[
-    [CallbackContext, LlmRequest], Optional[LlmResponse]
+BeforeModelCallback: TypeAlias = Union[
+    _SingleBeforeModelCallback,
+    list[_SingleBeforeModelCallback],
 ]
-AfterModelCallback: TypeAlias = Callable[
+
+_SingleAfterModelCallback: TypeAlias = Callable[
     [CallbackContext, LlmResponse],
-    Optional[LlmResponse],
+    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
 ]
+
+AfterModelCallback: TypeAlias = Union[
+    _SingleAfterModelCallback,
+    list[_SingleAfterModelCallback],
+]
+
 BeforeToolCallback: TypeAlias = Callable[
     [BaseTool, dict[str, Any], ToolContext],
     Union[Awaitable[Optional[dict]], Optional[dict]],
@@ -173,7 +185,11 @@ class LlmAgent(BaseAgent):
 
   # Callbacks - Start
   before_model_callback: Optional[BeforeModelCallback] = None
-  """Called before calling the LLM.
+  """Callback or list of callbacks to be called before calling the LLM.
+
+  When a list of callbacks is provided, the callbacks will be called in the
+  order they are listed until a callback does not return None.
+
   Args:
     callback_context: CallbackContext,
     llm_request: LlmRequest, The raw model request. Callback can mutate the
@@ -184,7 +200,10 @@ class LlmAgent(BaseAgent):
     skipped and the provided content will be returned to user.
   """
   after_model_callback: Optional[AfterModelCallback] = None
-  """Called after calling LLM.
+  """Callback or list of callbacks to be called after calling the LLM.
+
+  When a list of callbacks is provided, the callbacks will be called in the
+  order they are listed until a callback does not return None.
 
   Args:
     callback_context: CallbackContext,
@@ -284,6 +303,32 @@ def canonical_tools(self) -> list[BaseTool]:
     """
     return [_convert_tool_union_to_tool(tool) for tool in self.tools]
 
+  @property
+  def canonical_before_model_callbacks(
+      self,
+  ) -> list[_SingleBeforeModelCallback]:
+    """The resolved self.before_model_callback field as a list of _SingleBeforeModelCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.before_model_callback:
+      return []
+    if isinstance(self.before_model_callback, list):
+      return self.before_model_callback
+    return [self.before_model_callback]
+
+  @property
+  def canonical_after_model_callbacks(self) -> list[_SingleAfterModelCallback]:
+    """The resolved self.after_model_callback field as a list of _SingleAfterModelCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.after_model_callback:
+      return []
+    if isinstance(self.after_model_callback, list):
+      return self.after_model_callback
+    return [self.after_model_callback]
+
   @property
   def _llm_flow(self) -> BaseLlmFlow:
     if (

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

@@ -16,6 +16,7 @@
 
 from abc import ABC
 import asyncio
+import inspect
 import logging
 from typing import AsyncGenerator
 from typing import cast
@@ -192,14 +193,15 @@ async def _receive_from_model(
     """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.
+
+      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:
@@ -447,7 +449,7 @@ async def _call_llm_async(
       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
@@ -460,7 +462,7 @@ async def _call_llm_async(
         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
@@ -489,14 +491,14 @@ async def _call_llm_async(
               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,
@@ -508,17 +510,23 @@ def _handle_before_model_callback(
     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,
@@ -530,15 +538,21 @@ def _handle_after_model_callback(
     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,

src/google/adk/models/lite_llm.py

@@ -611,7 +611,7 @@ async def generate_content_async(
     """
 
     self._maybe_append_user_content(llm_request)
-    logger.info(_build_request_log(llm_request))
+    logger.debug(_build_request_log(llm_request))
 
     messages, tools = _get_completion_inputs(llm_request)
 

src/google/adk/tools/agent_tool.py

@@ -146,11 +146,12 @@ async def run_async(
 
     if runner.artifact_service:
       # Forward all artifacts to parent session.
-      async for artifact_name in runner.artifact_service.list_artifact_keys(
+      artifact_names = await runner.artifact_service.list_artifact_keys(
           app_name=session.app_name,
           user_id=session.user_id,
           session_id=session.id,
-      ):
+      )
+      for artifact_name in artifact_names:
         if artifact := await runner.artifact_service.load_artifact(
             app_name=session.app_name,
             user_id=session.user_id,

src/google/adk/tools/application_integration_tool/application_integration_toolset.py

@@ -76,7 +76,7 @@ def __init__(
       project: str,
       location: str,
       integration: Optional[str] = None,
-      trigger: Optional[str] = None,
+      triggers: Optional[List[str]] = None,
       connection: Optional[str] = None,
       entity_operations: Optional[str] = None,
       actions: Optional[str] = None,
@@ -98,7 +98,7 @@ def __init__(
         project="test-project",
         location="us-central1"
         integration="test-integration",
-        trigger="api_trigger/test_trigger",
+        triggers=["api_trigger/test_trigger"],
         service_account_credentials={...},
     )
 
@@ -130,7 +130,7 @@ def __init__(
         project: The GCP project ID.
         location: The GCP location.
         integration: The integration name.
-        trigger: The trigger name.
+        triggers: The list of trigger names in the integration.
         connection: The connection name.
         entity_operations: The entity operations supported by the connection.
         actions: The actions supported by the connection.
@@ -149,7 +149,7 @@ def __init__(
     self.project = project
     self.location = location
     self.integration = integration
-    self.trigger = trigger
+    self.triggers = triggers
     self.connection = connection
     self.entity_operations = entity_operations
     self.actions = actions
@@ -162,14 +162,14 @@ def __init__(
         project,
         location,
         integration,
-        trigger,
+        triggers,
         connection,
         entity_operations,
         actions,
         service_account_json,
     )
     connection_details = {}
-    if integration and trigger:
+    if integration:
       spec = integration_client.get_openapi_spec_for_integration()
     elif connection and (entity_operations or actions):
       connections_client = ConnectionsClient(
@@ -210,7 +210,7 @@ def _parse_spec_to_tools(self, spec_dict, connection_details):
       )
       auth_scheme = HTTPBearer(bearerFormat="JWT")
 
-    if self.integration and self.trigger:
+    if self.integration:
       tools = OpenAPIToolset(
           spec_dict=spec_dict,
           auth_credential=auth_credential,