google-adk 0.4.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

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/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.
@@ -332,7 +333,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

@@ -53,16 +53,24 @@ 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 = _populate_values(raw_si, invocation_context)
+      si = raw_si
+      if not bypass_state_injection:
+        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)
+      raw_si, bypass_state_injection = await agent.canonical_instruction(
+          ReadonlyContext(invocation_context)
+      )
+      si = raw_si
+      if not bypass_state_injection:
+        si = await _populate_values(raw_si, invocation_context)
       llm_request.append_instructions([si])
 
     # Maintain async generator behavior
@@ -73,13 +81,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 +108,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 +128,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/flows/llm_flows/single_flow.py

@@ -25,7 +25,7 @@ from . import identity
 from . import instructions
 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,44 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import abc
+
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from typing import TYPE_CHECKING
 
 from pydantic import BaseModel
 from pydantic import Field
 
-from ..events.event import Event
-from ..sessions.session import Session
+from .memory_entry import MemoryEntry
 
-
-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.
-  """
-  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
-  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.
@@ -58,9 +58,13 @@ class BaseMemoryService(abc.ABC):
         session: The session to add.
     """
 
-  @abc.abstractmethod
-  def search_memory(
-      self, *, app_name: str, user_id: str, query: str
+  @abstractmethod
+  async def search_memory(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      query: str,
   ) -> SearchMemoryResponse:
     """Searches for sessions that match the query.
 

google/adk/memory/base_memory_service.py.orig

@@ -0,0 +1,76 @@
+# 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.
+
+import abc
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from ..events.event import Event
+from ..sessions.session import Session
+
+
+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.
+  """
+
+  session_id: str
+  events: list[Event]
+
+
+class SearchMemoryResponse(BaseModel):
+  """Represents the response from a memory search.
+
+  Attributes:
+      memories: A list of memory results matching the search query.
+  """
+
+  memories: list[MemoryResult] = Field(default_factory=list)
+
+
+class BaseMemoryService(abc.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):
+    """Adds a session to the memory service.
+
+    A session may be added multiple times during its lifetime.
+
+    Args:
+        session: The session to add.
+    """
+
+  @abc.abstractmethod
+  async def search_memory(
+      self, *, app_name: str, user_id: str, query: str
+  ) -> SearchMemoryResponse:
+    """Searches for sessions that match the query.
+
+    Args:
+        app_name: The name of the application.
+        user_id: The id of the user.
+        query: The query to search for.
+
+    Returns:
+        A SearchMemoryResponse containing the matching memories.
+    """

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."""
 
-  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
+  @override
+  async def add_session_to_memory(self, session: Session):
+    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
     ]
 
-  def search_memory(
+  @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.
+  """

google/adk/memory/vertex_ai_rag_memory_service.py

@@ -12,20 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+from __future__ import annotations
+
 from collections import OrderedDict
 import json
 import os
 import tempfile
+from typing import Optional
+from typing import TYPE_CHECKING
 
 from google.genai import types
 from typing_extensions import override
 from vertexai.preview import rag
 
-from ..events.event import Event
-from ..sessions.session import Session
+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
 
 
 class VertexAiRagMemoryService(BaseMemoryService):
@@ -33,8 +41,8 @@ class VertexAiRagMemoryService(BaseMemoryService):
 
   def __init__(
       self,
-      rag_corpus: str = None,
-      similarity_top_k: int = None,
+      rag_corpus: Optional[str] = None,
+      similarity_top_k: Optional[int] = None,
       vector_distance_threshold: float = 10,
   ):
     """Initializes a VertexAiRagMemoryService.
@@ -47,14 +55,16 @@ class VertexAiRagMemoryService(BaseMemoryService):
         vector_distance_threshold: Only returns contexts with vector distance
           smaller than the threshold..
     """
-    self.vertex_rag_store = types.VertexRagStore(
-        rag_resources=[rag.RagResource(rag_corpus=rag_corpus)],
+    self._vertex_rag_store = types.VertexRagStore(
+        rag_resources=[
+            types.VertexRagStoreRagResource(rag_corpus=rag_corpus),
+        ],
         similarity_top_k=similarity_top_k,
         vector_distance_threshold=vector_distance_threshold,
     )
 
   @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:
@@ -79,7 +89,11 @@ class VertexAiRagMemoryService(BaseMemoryService):
       output_string = "\n".join(output_lines)
       temp_file.write(output_string)
       temp_file_path = temp_file.name
-    for rag_resource in self.vertex_rag_store.rag_resources:
+
+    if not self._vertex_rag_store.rag_resources:
+      raise ValueError("Rag resources must be set.")
+
+    for rag_resource in self._vertex_rag_store.rag_resources:
       rag.upload_file(
           corpus_name=rag_resource.rag_corpus,
           path=temp_file_path,
@@ -91,16 +105,18 @@ 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."""
+    from ..events.event import Event
+
     response = rag.retrieval_query(
         text=query,
-        rag_resources=self.vertex_rag_store.rag_resources,
-        rag_corpora=self.vertex_rag_store.rag_corpora,
-        similarity_top_k=self.vertex_rag_store.similarity_top_k,
-        vector_distance_threshold=self.vertex_rag_store.vector_distance_threshold,
+        rag_resources=self._vertex_rag_store.rag_resources,
+        rag_corpora=self._vertex_rag_store.rag_corpora,
+        similarity_top_k=self._vertex_rag_store.similarity_top_k,
+        vector_distance_threshold=self._vertex_rag_store.vector_distance_threshold,
     )
 
     memory_results = []
@@ -144,9 +160,16 @@ class VertexAiRagMemoryService(BaseMemoryService):
     for session_id, event_lists in session_events_map.items():
       for events in _merge_event_lists(event_lists):
         sorted_events = sorted(events, key=lambda e: e.timestamp)
-        memory_results.append(
-            MemoryResult(session_id=session_id, events=sorted_events)
-        )
+
+        memory_results.extend([
+            MemoryEntry(
+                author=event.author,
+                content=event.content,
+                timestamp=_utils.format_timestamp(event.timestamp),
+            )
+            for event in sorted_events
+            if event.content
+        ])
     return SearchMemoryResponse(memories=memory_results)