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

google/adk/flows/llm_flows/base_llm_flow.py

@@ -25,10 +25,12 @@ from typing import TYPE_CHECKING
 
 from websockets.exceptions import ConnectionClosedOK
 
+from . import functions
 from ...agents.base_agent import BaseAgent
 from ...agents.callback_context import CallbackContext
 from ...agents.invocation_context import InvocationContext
 from ...agents.live_request_queue import LiveRequestQueue
+from ...agents.readonly_context import ReadonlyContext
 from ...agents.run_config import StreamingMode
 from ...agents.transcription_entry import TranscriptionEntry
 from ...events.event import Event
@@ -39,7 +41,6 @@ from ...telemetry import trace_call_llm
 from ...telemetry import trace_send_data
 from ...telemetry import tracer
 from ...tools.tool_context import ToolContext
-from . import functions
 
 if TYPE_CHECKING:
   from ...agents.llm_agent import LlmAgent
@@ -47,7 +48,7 @@ if TYPE_CHECKING:
   from ._base_llm_processor import BaseLlmRequestProcessor
   from ._base_llm_processor import BaseLlmResponseProcessor
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class BaseLlmFlow(ABC):
@@ -88,7 +89,12 @@ class BaseLlmFlow(ABC):
           if invocation_context.transcription_cache:
             from . import audio_transcriber
 
-            audio_transcriber = audio_transcriber.AudioTranscriber()
+            audio_transcriber = audio_transcriber.AudioTranscriber(
+                init_client=True
+                if invocation_context.run_config.input_audio_transcription
+                is None
+                else False
+            )
             contents = audio_transcriber.transcribe_file(invocation_context)
             logger.debug('Sending history to model: %s', contents)
             await llm_connection.send_history(contents)
@@ -129,6 +135,18 @@ class BaseLlmFlow(ABC):
             # cancel the tasks that belongs to the closed connection.
             send_task.cancel()
             await llm_connection.close()
+          if (
+              event.content
+              and event.content.parts
+              and event.content.parts[0].function_response
+              and event.content.parts[0].function_response.name
+              == 'task_completed'
+          ):
+            # this is used for sequential agent to signal the end of the agent.
+            await asyncio.sleep(1)
+            # cancel the tasks that belongs to the closed connection.
+            send_task.cancel()
+            return
       finally:
         # Clean up
         if not send_task.done():
@@ -176,9 +194,12 @@ class BaseLlmFlow(ABC):
         # Cache audio data here for transcription
         if not invocation_context.transcription_cache:
           invocation_context.transcription_cache = []
-        invocation_context.transcription_cache.append(
-            TranscriptionEntry(role='user', data=live_request.blob)
-        )
+        if not invocation_context.run_config.input_audio_transcription:
+          # if the live model's input transcription is not enabled, then
+          # we use our onwn audio transcriber to achieve that.
+          invocation_context.transcription_cache.append(
+              TranscriptionEntry(role='user', data=live_request.blob)
+          )
         await llm_connection.send_realtime(live_request.blob)
       if live_request.content:
         await llm_connection.send_content(live_request.content)
@@ -191,14 +212,22 @@ class BaseLlmFlow(ABC):
       llm_request: LlmRequest,
   ) -> AsyncGenerator[Event, None]:
     """Receive data from model and process events using BaseLlmConnection."""
-    def get_author(llm_response):
+
+    def get_author_for_event(llm_response):
       """Get the author of the event.
 
       When the model returns transcription, the author is "user". Otherwise, the
-      author is the agent.
+      author is the agent name(not 'model').
+
+      Args:
+        llm_response: The LLM response from the LLM call.
       """
-      if llm_response and llm_response.content and llm_response.content.role == "user":
-        return "user"
+      if (
+          llm_response
+          and llm_response.content
+          and llm_response.content.role == 'user'
+      ):
+        return 'user'
       else:
         return invocation_context.agent.name
 
@@ -209,7 +238,7 @@ class BaseLlmFlow(ABC):
           model_response_event = Event(
               id=Event.new_id(),
               invocation_id=invocation_context.invocation_id,
-              author=get_author(llm_response),
+              author=get_author_for_event(llm_response),
           )
           async for event in self._postprocess_live(
               invocation_context,
@@ -220,13 +249,20 @@ class BaseLlmFlow(ABC):
             if (
                 event.content
                 and event.content.parts
-                and event.content.parts[0].text
+                and event.content.parts[0].inline_data is None
                 and not event.partial
             ):
+              # This can be either user data or transcription data.
+              # when output transcription enabled, it will contain model's
+              # transcription.
+              # when input transcription enabled, it will contain user
+              # transcription.
               if not invocation_context.transcription_cache:
                 invocation_context.transcription_cache = []
               invocation_context.transcription_cache.append(
-                  TranscriptionEntry(role='model', data=event.content)
+                  TranscriptionEntry(
+                      role=event.content.role, data=event.content
+                  )
               )
             yield event
         # Give opportunity for other tasks to run.
@@ -261,6 +297,7 @@ 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,
@@ -272,8 +309,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()
+        # Update the mutable event id to avoid conflict
+        model_response_event.id = Event.new_id()
         yield event
 
   async def _preprocess_async(
@@ -291,7 +328,9 @@ class BaseLlmFlow(ABC):
         yield event
 
     # Run processors for tools.
-    for tool in agent.canonical_tools:
+    for tool in await agent.canonical_tools(
+        ReadonlyContext(invocation_context)
+    ):
       tool_context = ToolContext(invocation_context)
       await tool.process_llm_request(
           tool_context=tool_context, llm_request=llm_request
@@ -433,14 +472,12 @@ class BaseLlmFlow(ABC):
           yield event
 
   def _get_agent_to_run(
-      self, invocation_context: InvocationContext, transfer_to_agent
+      self, invocation_context: InvocationContext, agent_name: str
   ) -> BaseAgent:
     root_agent = invocation_context.agent.root_agent
-    agent_to_run = root_agent.find_agent(transfer_to_agent)
+    agent_to_run = root_agent.find_agent(agent_name)
     if not agent_to_run:
-      raise ValueError(
-          f'Agent {transfer_to_agent} not found in the agent tree.'
-      )
+      raise ValueError(f'Agent {agent_name} not found in the agent tree.')
     return agent_to_run
 
   async def _call_llm_async(

google/adk/flows/llm_flows/contents.py

@@ -15,7 +15,9 @@
 from __future__ import annotations
 
 import copy
-from typing import AsyncGenerator, Generator, Optional
+from typing import AsyncGenerator
+from typing import Generator
+from typing import Optional
 
 from google.genai import types
 from typing_extensions import override

google/adk/flows/llm_flows/functions.py

@@ -41,7 +41,7 @@ from ...tools.tool_context import ToolContext
 AF_FUNCTION_CALL_ID_PREFIX = 'adk-'
 REQUEST_EUC_FUNCTION_CALL_NAME = 'adk_request_credential'
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 def generate_client_function_call_id() -> str:
@@ -106,7 +106,7 @@ def generate_auth_event(
         args=AuthToolArguments(
             function_call_id=function_call_id,
             auth_config=auth_config,
-        ).model_dump(exclude_none=True),
+        ).model_dump(exclude_none=True, by_alias=True),
     )
     request_euc_function_call.id = generate_client_function_call_id()
     long_running_tool_ids.add(request_euc_function_call.id)
@@ -153,22 +153,22 @@ async def handle_function_calls_async(
     function_args = function_call.args or {}
     function_response: Optional[dict] = None
 
-    # before_tool_callback (sync or async)
-    if agent.before_tool_callback:
-      function_response = agent.before_tool_callback(
+    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
       )
 
-    # after_tool_callback (sync or async)
-    if agent.after_tool_callback:
-      altered_function_response = agent.after_tool_callback(
+    for callback in agent.canonical_after_tool_callbacks:
+      altered_function_response = callback(
           tool=tool,
           args=function_args,
           tool_context=tool_context,
@@ -178,6 +178,7 @@ async def handle_function_calls_async(
         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.

google/adk/flows/llm_flows/instructions.py

@@ -26,6 +26,7 @@ from typing_extensions import override
 from ...agents.readonly_context import ReadonlyContext
 from ...events.event import Event
 from ...sessions.state import State
+from ...utils import instructions_utils
 from ._base_llm_processor import BaseLlmRequestProcessor
 
 if TYPE_CHECKING:
@@ -53,16 +54,28 @@ class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
     if (
         isinstance(root_agent, LlmAgent) and root_agent.global_instruction
     ):  # not empty str
-      raw_si = root_agent.canonical_global_instruction(
-          ReadonlyContext(invocation_context)
+      raw_si, bypass_state_injection = (
+          await root_agent.canonical_global_instruction(
+              ReadonlyContext(invocation_context)
+          )
       )
-      si = await _populate_values(raw_si, invocation_context)
+      si = raw_si
+      if not bypass_state_injection:
+        si = await instructions_utils.inject_session_state(
+            raw_si, ReadonlyContext(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 = await _populate_values(raw_si, invocation_context)
+      raw_si, bypass_state_injection = await agent.canonical_instruction(
+          ReadonlyContext(invocation_context)
+      )
+      si = raw_si
+      if not bypass_state_injection:
+        si = await instructions_utils.inject_session_state(
+            raw_si, ReadonlyContext(invocation_context)
+        )
       llm_request.append_instructions([si])
 
     # Maintain async generator behavior
@@ -71,78 +84,3 @@ class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
 
 
 request_processor = _InstructionsLlmRequestProcessor()
-
-
-async def _populate_values(
-    instruction_template: str,
-    context: InvocationContext,
-) -> str:
-  """Populates values in the instruction template, e.g. state, artifact, etc."""
-
-  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('?'):
-      optional = True
-      var_name = var_name.removesuffix('?')
-    if var_name.startswith('artifact.'):
-      var_name = var_name.removeprefix('artifact.')
-      if context.artifact_service is None:
-        raise ValueError('Artifact service is not initialized.')
-      artifact = await context.artifact_service.load_artifact(
-          app_name=context.session.app_name,
-          user_id=context.session.user_id,
-          session_id=context.session.id,
-          filename=var_name,
-      )
-      if not var_name:
-        raise KeyError(f'Artifact {var_name} not found.')
-      return str(artifact)
-    else:
-      if not _is_valid_state_name(var_name):
-        return match.group()
-      if var_name in context.session.state:
-        return str(context.session.state[var_name])
-      else:
-        if optional:
-          return ''
-        else:
-          raise KeyError(f'Context variable not found: `{var_name}`.')
-
-  return await _async_sub(r'{+[^{}]*}+', _replace_match, instruction_template)
-
-
-def _is_valid_state_name(var_name):
-  """Checks if the variable name is a valid state name.
-
-  Valid state is either:
-    - Valid identifier
-    - <Valid prefix>:<Valid identifier>
-  All the others will just return as it is.
-
-  Args:
-    var_name: The variable name to check.
-
-  Returns:
-    True if the variable name is a valid state name, False otherwise.
-  """
-  parts = var_name.split(':')
-  if len(parts) == 1:
-    return var_name.isidentifier()
-
-  if len(parts) == 2:
-    prefixes = [State.APP_PREFIX, State.USER_PREFIX, State.TEMP_PREFIX]
-    if (parts[0] + ':') in prefixes:
-      return parts[1].isidentifier()
-  return False

google/adk/flows/llm_flows/single_flow.py

@@ -16,16 +16,16 @@
 
 import logging
 
-from ...auth import auth_preprocessor
 from . import _code_execution
 from . import _nl_planning
 from . import basic
 from . import contents
 from . import identity
 from . import instructions
+from ...auth import auth_preprocessor
 from .base_llm_flow import BaseLlmFlow
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class SingleFlow(BaseLlmFlow):

google/adk/memory/__init__.py

@@ -16,7 +16,7 @@ import logging
 from .base_memory_service import BaseMemoryService
 from .in_memory_memory_service import InMemoryMemoryService
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 __all__ = [
     'BaseMemoryService',

google/adk/memory/_utils.py

@@ -0,0 +1,23 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+
+from __future__ import annotations
+
+from datetime import datetime
+
+
+def format_timestamp(timestamp: float) -> str:
+  """Formats the timestamp of the memory entry."""
+  return datetime.fromtimestamp(timestamp).isoformat()

google/adk/memory/base_memory_service.py

@@ -12,46 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import abc
 
-from pydantic import BaseModel
-from pydantic import Field
+from __future__ import annotations
 
-from ..events.event import Event
-from ..sessions.session import Session
+from abc import ABC
+from abc import abstractmethod
+from typing import TYPE_CHECKING
 
+from pydantic import BaseModel
+from pydantic import Field
 
-class MemoryResult(BaseModel):
-  """Represents a single memory retrieval result.
-
-  Attributes:
-      session_id: The session id associated with the memory.
-      events: A list of events in the session.
-  """
+from .memory_entry import MemoryEntry
 
-  session_id: str
-  events: list[Event]
+if TYPE_CHECKING:
+  from ..sessions.session import Session
 
 
 class SearchMemoryResponse(BaseModel):
   """Represents the response from a memory search.
 
   Attributes:
-      memories: A list of memory results matching the search query.
+      memories: A list of memory entries that relate to the search query.
   """
 
-  memories: list[MemoryResult] = Field(default_factory=list)
+  memories: list[MemoryEntry] = Field(default_factory=list)
 
 
-class BaseMemoryService(abc.ABC):
+class BaseMemoryService(ABC):
   """Base class for memory services.
 
   The service provides functionalities to ingest sessions into memory so that
   the memory can be used for user queries.
   """
 
-  @abc.abstractmethod
-  async def add_session_to_memory(self, session: Session):
+  @abstractmethod
+  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.
@@ -60,9 +58,13 @@ class BaseMemoryService(abc.ABC):
         session: The session to add.
     """
 
-  @abc.abstractmethod
+  @abstractmethod
   async def search_memory(
-      self, *, app_name: str, user_id: str, query: str
+      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

@@ -12,11 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ..events.event import Event
-from ..sessions.session import Session
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from . import _utils
 from .base_memory_service import BaseMemoryService
-from .base_memory_service import MemoryResult
 from .base_memory_service import SearchMemoryResponse
+from .memory_entry import MemoryEntry
+
+if TYPE_CHECKING:
+  from ..events.event import Event
+  from ..sessions.session import Session
+
+
+def _user_key(app_name: str, user_id: str):
+  return f'{app_name}/{user_id}'
+
+
+def _extract_words_lower(text: str) -> set[str]:
+  """Extracts words from a string and converts them to lowercase."""
+  return set([word.lower() for word in re.findall(r'[A-Za-z]+', text)])
 
 
 class InMemoryMemoryService(BaseMemoryService):
@@ -26,37 +46,49 @@ class InMemoryMemoryService(BaseMemoryService):
   """
 
   def __init__(self):
-    self.session_events: dict[str, list[Event]] = {}
-    """keys are app_name/user_id/session_id"""
+    self._session_events: dict[str, dict[str, list[Event]]] = {}
+    """Keys are app_name/user_id, session_id. Values are session event lists."""
 
+  @override
   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
+    user_key = _user_key(session.app_name, session.user_id)
+    self._session_events[user_key] = self._session_events.get(
+        _user_key(session.app_name, session.user_id), {}
+    )
+    self._session_events[user_key][session.id] = [
+        event
+        for event in session.events
+        if event.content and event.content.parts
     ]
 
+  @override
   async def search_memory(
       self, *, app_name: str, user_id: str, query: str
   ) -> SearchMemoryResponse:
-    """Prototyping purpose only."""
-    keywords = set(query.lower().split())
+    user_key = _user_key(app_name, user_id)
+    if user_key not in self._session_events:
+      return SearchMemoryResponse()
+
+    words_in_query = set(query.lower().split())
     response = SearchMemoryResponse()
-    for key, events in self.session_events.items():
-      if not key.startswith(f'{app_name}/{user_id}/'):
-        continue
-      matched_events = []
-      for event in events:
+
+    for session_events in self._session_events[user_key].values():
+      for event in session_events:
         if not event.content or not event.content.parts:
           continue
-        parts = event.content.parts
-        text = '\n'.join([part.text for part in parts if part.text]).lower()
-        for keyword in keywords:
-          if keyword in text:
-            matched_events.append(event)
-            break
-      if matched_events:
-        session_id = key.split('/')[-1]
-        response.memories.append(
-            MemoryResult(session_id=session_id, events=matched_events)
+        words_in_event = _extract_words_lower(
+            ' '.join([part.text for part in event.content.parts if part.text])
         )
+        if not words_in_event:
+          continue
+
+        if any(query_word in words_in_event for query_word in words_in_query):
+          response.memories.append(
+              MemoryEntry(
+                  content=event.content,
+                  author=event.author,
+                  timestamp=_utils.format_timestamp(event.timestamp),
+              )
+          )
+
     return response

google/adk/memory/memory_entry.py

@@ -0,0 +1,37 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+
+from __future__ import annotations
+
+from typing import Optional
+
+from google.genai import types
+from pydantic import BaseModel
+
+
+class MemoryEntry(BaseModel):
+  """Represent one memory entry."""
+
+  content: types.Content
+  """The main content of the memory."""
+
+  author: Optional[str] = None
+  """The author of the memory."""
+
+  timestamp: Optional[str] = None
+  """The timestamp when the original content of this memory happened.
+
+  This string will be forwarded to LLM. Preferred format is ISO 8601 format.
+  """