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

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/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."""
 
+  @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.
+  """

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,8 +55,10 @@ 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,
     )
@@ -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,
@@ -95,12 +109,14 @@ class VertexAiRagMemoryService(BaseMemoryService):
       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)
 
 

google/adk/models/anthropic_llm.py

@@ -24,8 +24,9 @@ from typing import AsyncGenerator
 from typing import Generator
 from typing import Iterable
 from typing import Literal
-from typing import Optional, Union
+from typing import Optional
 from typing import TYPE_CHECKING
+from typing import Union
 
 from anthropic import AnthropicVertex
 from anthropic import NOT_GIVEN
@@ -42,7 +43,7 @@ if TYPE_CHECKING:
 
 __all__ = ["Claude"]
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 MAX_TOKEN = 1024
 
@@ -140,15 +141,15 @@ def message_to_generate_content_response(
           role="model",
           parts=[content_block_to_part(cb) for cb in message.content],
       ),
+      usage_metadata=types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=message.usage.input_tokens,
+          candidates_token_count=message.usage.output_tokens,
+          total_token_count=(
+              message.usage.input_tokens + message.usage.output_tokens
+          ),
+      ),
       # TODO: Deal with these later.
       # finish_reason=to_google_genai_finish_reason(message.stop_reason),
-      # usage_metadata=types.GenerateContentResponseUsageMetadata(
-      #     prompt_token_count=message.usage.input_tokens,
-      #     candidates_token_count=message.usage.output_tokens,
-      #     total_token_count=(
-      #         message.usage.input_tokens + message.usage.output_tokens
-      #     ),
-      # ),
   )
 
 
@@ -196,6 +197,12 @@ def function_declaration_to_tool_param(
 
 
 class Claude(BaseLlm):
+  """ "Integration with Claude models served from Vertex AI.
+
+  Attributes:
+    model: The name of the Claude model.
+  """
+
   model: str = "claude-3-5-sonnet-v2@20241022"
 
   @staticmethod

google/adk/models/gemini_llm_connection.py

@@ -21,7 +21,7 @@ from google.genai import types
 from .base_llm_connection import BaseLlmConnection
 from .llm_response import LlmResponse
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class GeminiLlmConnection(BaseLlmConnection):
@@ -149,16 +149,16 @@ class GeminiLlmConnection(BaseLlmConnection):
             message.server_content.input_transcription
             and message.server_content.input_transcription.text
         ):
-            user_text = message.server_content.input_transcription.text
-            parts = [
-                types.Part.from_text(
-                    text=user_text,
-                )
-            ]
-            llm_response = LlmResponse(
-                content=types.Content(role='user', parts=parts)
-            )
-            yield llm_response
+          user_text = message.server_content.input_transcription.text
+          parts = [
+              types.Part.from_text(
+                  text=user_text,
+              )
+          ]
+          llm_response = LlmResponse(
+              content=types.Content(role='user', parts=parts)
+          )
+          yield llm_response
         if (
             message.server_content.output_transcription
             and message.server_content.output_transcription.text

google/adk/models/google_llm.py

@@ -11,6 +11,8 @@
 # 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
 
 import contextlib
@@ -34,7 +36,7 @@ from .llm_response import LlmResponse
 if TYPE_CHECKING:
   from .llm_request import LlmRequest
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 _NEW_LINE = '\n'
 _EXCLUDED_PART_FIELD = {'inline_data': {'data'}}
@@ -121,6 +123,7 @@ class Gemini(BaseLlm):
               content=types.ModelContent(
                   parts=[types.Part.from_text(text=text)],
               ),
+              usage_metadata=llm_response.usage_metadata,
           )
           text = ''
         yield llm_response
@@ -174,9 +177,13 @@ class Gemini(BaseLlm):
   @cached_property
   def _live_api_client(self) -> Client:
     if self._api_backend == 'vertex':
+      # use beta version for vertex api
+      api_version = 'v1beta1'
       # use default api version for vertex
       return Client(
-          http_options=types.HttpOptions(headers=self._tracking_headers)
+          http_options=types.HttpOptions(
+              headers=self._tracking_headers, api_version=api_version
+          )
       )
     else:
       # use v1alpha for ml_dev