From 0450d3aa4b58fe5e8594d24fbd777cf517de9f94 Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 12:15:49 +0530 Subject: [PATCH 1/9] feat: Add OpenMemoryService support and update dependencies - Introduced OpenMemoryService and OpenMemoryServiceConfig to the memory module. - Updated pyproject.toml to include httpx as a dependency with version constraints. - Added import error handling for OpenMemoryService, providing guidance for installation if httpx is not present. This enhancement allows for improved memory management capabilities within the application. --- contributing/samples/open_memory/README.md | 152 +++++++ contributing/samples/open_memory/__init__.py | 16 + contributing/samples/open_memory/agent.py | 39 ++ contributing/samples/open_memory/main.py | 143 +++++++ pyproject.toml | 6 + src/google/adk/memory/__init__.py | 15 + src/google/adk/memory/open_memory_service.py | 392 ++++++++++++++++++ .../memory/test_open_memory_service.py | 363 ++++++++++++++++ 8 files changed, 1126 insertions(+) create mode 100644 contributing/samples/open_memory/README.md create mode 100644 contributing/samples/open_memory/__init__.py create mode 100644 contributing/samples/open_memory/agent.py create mode 100644 contributing/samples/open_memory/main.py create mode 100644 src/google/adk/memory/open_memory_service.py create mode 100644 tests/unittests/memory/test_open_memory_service.py diff --git a/contributing/samples/open_memory/README.md b/contributing/samples/open_memory/README.md new file mode 100644 index 0000000000..0abaac8267 --- /dev/null +++ b/contributing/samples/open_memory/README.md @@ -0,0 +1,152 @@ +# OpenMemory Sample + +This sample demonstrates how to use OpenMemory as a self-hosted memory backend +for ADK agents. + +## Prerequisites + +- Python 3.9+ (Python 3.11+ recommended) +- Docker (for running OpenMemory) +- ADK installed with dependencies + +## Setup + +### 1. Start OpenMemory Server + +Start OpenMemory using Docker: + +```bash +docker run -p 3000:3000 cavira/openmemory +``` + +Or use the production network build: + +```bash +docker run -p 3000:3000 cavira/openmemory:production +``` + +Verify it's running: + +```bash +curl http://localhost:3000/health +``` + +### 2. Install Dependencies + +Install ADK with OpenMemory support: + +```bash +pip install google-adk[openmemory] +``` + +This installs `httpx` for making HTTP requests to the OpenMemory API. + +### 3. Configure Environment Variables + +Create a `.env` file in this directory (optional): + +```bash +# Required: Google API key for the agent +GOOGLE_API_KEY=your-google-api-key + +# Optional: OpenMemory base URL (defaults to localhost:3000) +OPENMEMORY_BASE_URL=http://localhost:3000 + +# Optional: API key if your OpenMemory instance requires authentication +# OPENMEMORY_API_KEY=your-api-key +``` + +**Note:** API key is only needed if your OpenMemory server is configured with authentication. + +## Usage + +### Basic Usage + +```python +from google.adk.memory import OpenMemoryService +from google.adk.runners import InMemoryRunner + +# Create OpenMemory service with defaults +memory_service = OpenMemoryService( + base_url="http://localhost:3000" +) + +# Use with runner +runner = InMemoryRunner( + app_name="my_app", + agent=my_agent, + memory_service=memory_service +) +``` + +### Advanced Configuration + +```python +from google.adk.memory import OpenMemoryService, OpenMemoryServiceConfig + +# Custom configuration +config = OpenMemoryServiceConfig( + search_top_k=20, # Retrieve more memories per query + timeout=10.0, # Faster timeout for production + user_content_salience=0.9, # Higher importance for user messages + model_content_salience=0.75, # Medium importance for model responses + enable_metadata_tags=True # Use tags for filtering +) + +memory_service = OpenMemoryService( + base_url="http://localhost:3000", + api_key="your-api-key", + config=config +) +``` + +## Running the Sample + +```bash +cd contributing/samples/open_memory +python main.py +``` + +## Expected Output + +``` +----Session to create memory: ---------------------- +** User says: {'role': 'user', 'parts': [{'text': 'Hi'}]} +** model: Hello! How can I help you today? +... +Saving session to memory service... +------------------------------------------------------------------- +----Session to use memory: ---------------------- +** User says: {'role': 'user', 'parts': [{'text': 'What do I like to do?'}]} +** model: You like badminton. +... +------------------------------------------------------------------- +``` + +## Configuration Options + +### OpenMemoryServiceConfig + +- `search_top_k` (int, default: 10): Maximum memories to retrieve per search +- `timeout` (float, default: 30.0): HTTP request timeout in seconds +- `user_content_salience` (float, default: 0.8): Importance for user messages +- `model_content_salience` (float, default: 0.7): Importance for model responses +- `default_salience` (float, default: 0.6): Fallback importance value +- `enable_metadata_tags` (bool, default: True): Include session/app tags + +## Features + +OpenMemory provides: + +- **Multi-sector embeddings**: Factual, emotional, temporal, relational memory +- **Graceful decay curves**: Automatic reinforcement keeps relevant context sharp +- **Self-hosted**: Full data ownership, no vendor lock-in +- **High performance**: 2-3× faster than hosted alternatives +- **Cost-effective**: 6-10× cheaper than SaaS memory APIs + +## Learn More + +- [OpenMemory Documentation](https://openmemory.cavira.app/) +- [OpenMemory API Reference](https://openmemory.cavira.app/docs/api/add-memory) +- [ADK Memory Documentation](https://google.github.io/adk-docs) + diff --git a/contributing/samples/open_memory/__init__.py b/contributing/samples/open_memory/__init__.py new file mode 100644 index 0000000000..8ce90a27ba --- /dev/null +++ b/contributing/samples/open_memory/__init__.py @@ -0,0 +1,16 @@ +# 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 . import agent + diff --git a/contributing/samples/open_memory/agent.py b/contributing/samples/open_memory/agent.py new file mode 100644 index 0000000000..8a61a7db68 --- /dev/null +++ b/contributing/samples/open_memory/agent.py @@ -0,0 +1,39 @@ +# 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 datetime import datetime + +from google.adk import Agent +from google.adk.agents.callback_context import CallbackContext +from google.adk.tools.load_memory_tool import load_memory_tool +from google.adk.tools.preload_memory_tool import preload_memory_tool + + +def update_current_time(callback_context: CallbackContext): + callback_context.state['_time'] = datetime.now().isoformat() + + +root_agent = Agent( + model='gemini-2.5-flash', + name='open_memory_agent', + description='agent that has access to memory tools with OpenMemory.', + before_agent_callback=update_current_time, + instruction="""\ + You are an agent that help user answer questions. You have access to memory tools. + You can use the memory tools to look up the information in the memory. Current time: {_time} + """, + tools=[load_memory_tool, preload_memory_tool], +) + diff --git a/contributing/samples/open_memory/main.py b/contributing/samples/open_memory/main.py new file mode 100644 index 0000000000..822b861f2b --- /dev/null +++ b/contributing/samples/open_memory/main.py @@ -0,0 +1,143 @@ +# 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 asyncio +import os +from datetime import datetime +from datetime import timedelta +from typing import cast + +import agent +from dotenv import load_dotenv +from google.adk.cli.utils import logs +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.memory.open_memory_service import ( + OpenMemoryService, + OpenMemoryServiceConfig, +) +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.sessions.session import Session +from google.genai import types + +load_dotenv(override=True) +logs.log_to_tmp_folder() + + +async def main(): + app_name = 'my_app' + user_id_1 = 'user1' + + # Configure OpenMemory from environment variables or use defaults + openmemory_base_url = os.getenv('OPENMEMORY_BASE_URL', 'http://localhost:3000') + openmemory_api_key = os.getenv('OPENMEMORY_API_KEY') # Optional + + # Create OpenMemory service with custom configuration + config = OpenMemoryServiceConfig( + search_top_k=10, + user_content_salience=0.8, + model_content_salience=0.7, + enable_metadata_tags=True + ) + memory_service = OpenMemoryService( + base_url=openmemory_base_url, + api_key=openmemory_api_key, + config=config + ) + + runner = Runner( + app_name=app_name, + agent=agent.root_agent, + artifact_service=InMemoryArtifactService(), + session_service=InMemorySessionService(), + memory_service=memory_service, + ) + + async def run_prompt(session: Session, new_message: str) -> Session: + content = types.Content( + role='user', parts=[types.Part.from_text(text=new_message)] + ) + print('** User says:', content.model_dump(exclude_none=True)) + async for event in runner.run_async( + user_id=user_id_1, + session_id=session.id, + new_message=content, + ): + if not event.content or not event.content.parts: + continue + if event.content.parts[0].text: + print(f'** {event.author}: {event.content.parts[0].text}') + elif event.content.parts[0].function_call: + print( + f'** {event.author}: fc /' + f' {event.content.parts[0].function_call.name} /' + f' {event.content.parts[0].function_call.args}\n' + ) + elif event.content.parts[0].function_response: + print( + f'** {event.author}: fr /' + f' {event.content.parts[0].function_response.name} /' + f' {event.content.parts[0].function_response.response}\n' + ) + + return cast( + Session, + await runner.session_service.get_session( + app_name=app_name, user_id=user_id_1, session_id=session.id + ), + ) + + session_1 = await runner.session_service.create_session( + app_name=app_name, user_id=user_id_1 + ) + + print( + f'----Session to create memory: {session_1.id} ----------------------' + ) + session_1 = await run_prompt(session_1, 'Hi') + session_1 = await run_prompt(session_1, 'My name is Jack') + session_1 = await run_prompt(session_1, 'I like badminton.') + session_1 = await run_prompt( + session_1, + f'I ate a burger on {(datetime.now() - timedelta(days=1)).date()}.', + ) + session_1 = await run_prompt( + session_1, + f'I ate a banana on {(datetime.now() - timedelta(days=2)).date()}.', + ) + print('Saving session to memory service...') + if runner.memory_service: + await runner.memory_service.add_session_to_memory(session_1) + print('-------------------------------------------------------------------') + + session_2 = await runner.session_service.create_session( + app_name=app_name, user_id=user_id_1 + ) + print(f'----Session to use memory: {session_2.id} ----------------------') + session_2 = await run_prompt(session_2, 'Hi') + session_2 = await run_prompt(session_2, 'What do I like to do?') + # ** open_memory_agent: You like badminton. + session_2 = await run_prompt(session_2, 'When did I say that?') + # ** open_memory_agent: You said you liked badminton on ... + session_2 = await run_prompt(session_2, 'What did I eat yesterday?') + # ** open_memory_agent: You ate a burger yesterday... + print('-------------------------------------------------------------------') + + # Cleanup + await memory_service.close() + + +if __name__ == '__main__': + asyncio.run(main()) + diff --git a/pyproject.toml b/pyproject.toml index 0ba44779e3..3067207bb1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,6 +108,12 @@ eval = [ # go/keep-sorted end ] +openmemory = [ + # go/keep-sorted start + "httpx>=0.27.0, <1.0.0", + # go/keep-sorted end +] + test = [ # go/keep-sorted start "a2a-sdk>=0.3.0,<0.4.0;python_version>='3.10'", diff --git a/src/google/adk/memory/__init__.py b/src/google/adk/memory/__init__.py index 915d7e5178..161e80f125 100644 --- a/src/google/adk/memory/__init__.py +++ b/src/google/adk/memory/__init__.py @@ -35,3 +35,18 @@ ' VertexAiRagMemoryService please install it. If not, you can ignore this' ' warning.' ) + +try: + from .open_memory_service import ( + OpenMemoryService, + OpenMemoryServiceConfig, + ) + + __all__.append('OpenMemoryService') + __all__.append('OpenMemoryServiceConfig') +except ImportError: + logger.debug( + 'httpx is not installed. If you want to use the OpenMemoryService' + ' please install it via "pip install google-adk[openmemory]".' + ' If not, you can ignore this warning.' + ) diff --git a/src/google/adk/memory/open_memory_service.py b/src/google/adk/memory/open_memory_service.py new file mode 100644 index 0000000000..fe78356de9 --- /dev/null +++ b/src/google/adk/memory/open_memory_service.py @@ -0,0 +1,392 @@ +# 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 + +import logging +import re +from typing import Optional +from typing import TYPE_CHECKING + +import httpx +from google.genai import types +from pydantic import BaseModel +from pydantic import Field +from typing_extensions import override + +from . import _utils +from .base_memory_service import BaseMemoryService +from .base_memory_service import SearchMemoryResponse +from .memory_entry import MemoryEntry + +if TYPE_CHECKING: + from ..sessions.session import Session + +logger = logging.getLogger('google_adk.' + __name__) + +class OpenMemoryService(BaseMemoryService): + """Memory service implementation using OpenMemory. + See https://openmemory.cavira.app/ for more information. + + OpenMemory provides hierarchical memory decomposition with multi-sector + embeddings, graceful decay curves, and automatic reinforcement for AI agents. + + Implementation Note: + This service uses direct HTTP requests via httpx for all operations. + This ensures full control over the request format, particularly for + passing `user_id` as a top-level field (required for server-side + filtering in queries). + + Example: + ```python + from google.adk.memory import OpenMemoryService, OpenMemoryServiceConfig + + # Basic usage with defaults + memory_service = OpenMemoryService( + base_url="http://localhost:3000" + ) + + # Custom configuration + config = OpenMemoryServiceConfig( + search_top_k=20, + user_content_salience=0.9, + model_content_salience=0.75 + ) + memory_service = OpenMemoryService( + base_url="http://localhost:3000", + api_key="my-api-key", + config=config + ) + ``` + """ + + def __init__( + self, + base_url: str = "http://localhost:3000", + api_key: Optional[str] = None, + config: Optional[OpenMemoryServiceConfig] = None, + ): + """Initializes the OpenMemory service. + + Args: + base_url: Base URL of the OpenMemory instance (default: localhost:3000). + api_key: API key for authentication (optional, only if server requires). + config: OpenMemoryServiceConfig instance for customizing behavior. If + None, uses default configuration. + """ + self._base_url = base_url.rstrip('/') + self._api_key = api_key + self._config = config or OpenMemoryServiceConfig() + + def _determine_salience(self, author: Optional[str]) -> float: + """Determine salience value based on content author. + + Args: + author: The author of the content (e.g., 'user', 'model', 'system'). + + Returns: + Salience value between 0.0 and 1.0. + """ + if not author: + return self._config.default_salience + + author_lower = author.lower() + if author_lower == "user": + return self._config.user_content_salience + elif author_lower == "model": + return self._config.model_content_salience + else: + return self._config.default_salience + + def _prepare_memory_data( + self, event, content_text: str, session + ) -> dict: + """Prepare memory data structure for OpenMemory API. + + Embeds author and timestamp directly in the content string so they + are available during search without needing additional API calls. + This avoids N+1 query problem when searching. + + Args: + event: The event to create memory from. + content_text: Extracted text content. + session: The session containing the event. + + Returns: + Dictionary with memory data formatted for OpenMemory API. + """ + # Format timestamp for display + timestamp_str = None + if event.timestamp: + timestamp_str = _utils.format_timestamp(event.timestamp) + + # Embed author and timestamp in content for retrieval during search + # Format: [Author: user, Time: 2025-11-04T10:32:01] Content text + enriched_content = content_text + metadata_parts = [] + if event.author: + metadata_parts.append(f"Author: {event.author}") + if timestamp_str: + metadata_parts.append(f"Time: {timestamp_str}") + + if metadata_parts: + metadata_prefix = "[" + ", ".join(metadata_parts) + "] " + enriched_content = metadata_prefix + content_text + + # Store metadata for filtering and tracking + metadata = { + "app_name": session.app_name, + "user_id": session.user_id, + "session_id": session.id, + "event_id": event.id, + "invocation_id": event.invocation_id, + "author": event.author, + "timestamp": timestamp_str, + "source": "adk_session" + } + metadata = {k: v for k, v in metadata.items() if v is not None} + + memory_data = { + "content": enriched_content, + "metadata": metadata, + "salience": self._determine_salience(event.author) + } + + if self._config.enable_metadata_tags: + memory_data["tags"] = [ + f"session:{session.id}", + f"app:{session.app_name}", + f"author:{event.author}" if event.author else None + ] + # Remove None values + memory_data["tags"] = [t for t in memory_data["tags"] if t] + + return memory_data + + @override + async def add_session_to_memory(self, session: Session): + """Add a session's events to OpenMemory. + + Processes all events in the session, filters out empty content, + and adds meaningful memories to OpenMemory with appropriate metadata. + + Args: + session: The session containing events to add to memory. + """ + memories_added = 0 + + headers = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + + async with httpx.AsyncClient(timeout=self._config.timeout) as http_client: + for event in session.events: + content_text = _extract_text_from_event(event) + if not content_text: + continue + + memory_data = self._prepare_memory_data(event, content_text, session) + + try: + # Include user_id as separate field for database storage and filtering + payload = { + "content": memory_data["content"], + "tags": memory_data.get("tags", []), + "metadata": memory_data.get("metadata", {}), + "salience": memory_data.get("salience", 0.5), + "user_id": session.user_id # Separate field for DB column + } + + response = await http_client.post( + f"{self._base_url}/memory/add", + json=payload, + headers=headers + ) + response.raise_for_status() + + memories_added += 1 + logger.debug("Added memory for event %s", event.id) + except Exception as e: + logger.error("Failed to add memory for event %s: %s", event.id, e) + + logger.info( + "Added %d memories from session %s", memories_added, session.id + ) + + def _build_search_payload( + self, app_name: str, user_id: str, query: str + ) -> dict: + """Build search payload for OpenMemory query API. + + Args: + app_name: The application name to filter by. + user_id: The user ID to filter by. + query: The search query string. + + Returns: + Dictionary with query parameters formatted for HTTP API. + """ + payload = { + "query": query, + "k": self._config.search_top_k, # Backend expects 'k', not 'top_k' + "filter": {} + } + + # Always filter by user_id for multi-user isolation + payload["filter"]["user_id"] = user_id + + # Add tag-based filtering if enabled + if self._config.enable_metadata_tags: + payload["filter"]["tags"] = [f"app:{app_name}"] + + return payload + + def _convert_to_memory_entry(self, result: dict) -> Optional[MemoryEntry]: + """Convert OpenMemory result to MemoryEntry. + + Extracts author and timestamp from the enriched content format: + [Author: user, Time: 2025-11-04T10:32:01] Content text + + Args: + result: OpenMemory search result (match from query or full memory). + + Returns: + MemoryEntry or None if conversion fails. + """ + try: + raw_content = result["content"] + author = None + timestamp = None + clean_content = raw_content + + # Parse enriched content format: [Author: X, Time: Y] Content + match = re.match(r'^\[([^\]]+)\]\s+(.*)', raw_content, re.DOTALL) + if match: + metadata_str = match.group(1) + clean_content = match.group(2) + + # Extract author + author_match = re.search(r'Author:\s*([^,\]]+)', metadata_str) + if author_match: + author = author_match.group(1).strip() + + # Extract timestamp + time_match = re.search(r'Time:\s*([^,\]]+)', metadata_str) + if time_match: + timestamp = time_match.group(1).strip() + + # Create content with clean text (without metadata prefix) + content = types.Content(parts=[types.Part(text=clean_content)]) + + return MemoryEntry( + content=content, + author=author, + timestamp=timestamp + ) + except (KeyError, ValueError) as e: + logger.debug("Failed to convert result to MemoryEntry: %s", e) + return None + + @override + async def search_memory( + self, *, app_name: str, user_id: str, query: str + ) -> SearchMemoryResponse: + """Search for memories using OpenMemory's query API. + + Queries OpenMemory with the search string and filters results + by app_name and user_id using tags (if metadata tagging is enabled). + + Args: + app_name: The application name to filter memories by. + user_id: The user ID to filter memories by. + query: The search query string. + + Returns: + SearchMemoryResponse containing matching memories. + """ + try: + search_payload = self._build_search_payload(app_name, user_id, query) + + # Use direct HTTP call for query since SDK v0.2.0 filters don't work properly + # The SDK doesn't correctly pass filter parameters to the backend + memories = [] + + async with httpx.AsyncClient(timeout=self._config.timeout) as http_client: + # Query for matching memories + headers = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + + # Debug: Log the exact payload being sent + logger.debug("Query payload: %s", search_payload) + + response = await http_client.post( + f"{self._base_url}/memory/query", + json=search_payload, + headers=headers + ) + response.raise_for_status() + result = response.json() + + # Debug: Log response summary + logger.debug("Query returned %d matches", len(result.get("matches", []))) + + # Backend returns 'matches' with content already including author/timestamp + # No need for additional API calls - parse from enriched content format + for match in result.get("matches", []): + memory_entry = self._convert_to_memory_entry(match) + if memory_entry: + memories.append(memory_entry) + + logger.info("Found %d memories for query: '%s'", len(memories), query) + return SearchMemoryResponse(memories=memories) + + except Exception as e: + logger.error("Failed to search memories: %s", e) + return SearchMemoryResponse(memories=[]) + + +class OpenMemoryServiceConfig(BaseModel): + """Configuration for OpenMemory service behavior. + + Attributes: + search_top_k: Maximum number of memories to retrieve per search. + timeout: Request timeout in seconds. + user_content_salience: Salience for user-authored content (0.0-1.0). + model_content_salience: Salience for model-generated content (0.0-1.0). + default_salience: Default salience value for memories (0.0-1.0). + enable_metadata_tags: Include session/app tags in memories. + """ + + search_top_k: int = Field(default=10, ge=1, le=100) + timeout: float = Field(default=30.0, gt=0.0) + user_content_salience: float = Field(default=0.8, ge=0.0, le=1.0) + model_content_salience: float = Field(default=0.7, ge=0.0, le=1.0) + default_salience: float = Field(default=0.6, ge=0.0, le=1.0) + enable_metadata_tags: bool = Field(default=True) + +def _extract_text_from_event(event) -> str: + """Extracts text content from an event's content parts. + + Args: + event: The event to extract text from. + + Returns: + Combined text from all text parts, or empty string if none found. + """ + if not event.content or not event.content.parts: + return '' + + text_parts = [part.text for part in event.content.parts if part.text] + return ' '.join(text_parts) diff --git a/tests/unittests/memory/test_open_memory_service.py b/tests/unittests/memory/test_open_memory_service.py new file mode 100644 index 0000000000..88ef1cff39 --- /dev/null +++ b/tests/unittests/memory/test_open_memory_service.py @@ -0,0 +1,363 @@ +# 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 unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.events.event import Event +from google.adk.memory.open_memory_service import ( + OpenMemoryService, + OpenMemoryServiceConfig, +) +from google.adk.sessions.session import Session +from google.genai import types +import pytest + +MOCK_APP_NAME = 'test-app' +MOCK_USER_ID = 'test-user' +MOCK_SESSION_ID = 'session-1' + +MOCK_SESSION = Session( + app_name=MOCK_APP_NAME, + user_id=MOCK_USER_ID, + id=MOCK_SESSION_ID, + last_update_time=1000, + events=[ + Event( + id='event-1', + invocation_id='inv-1', + author='user', + timestamp=12345, + content=types.Content(parts=[types.Part(text='Hello, I like Python.')]), + ), + Event( + id='event-2', + invocation_id='inv-2', + author='model', + timestamp=12346, + content=types.Content( + parts=[types.Part(text='Python is a great programming language.')] + ), + ), + # Empty event, should be ignored + Event( + id='event-3', + invocation_id='inv-3', + author='user', + timestamp=12347, + ), + # Function call event, should be ignored + Event( + id='event-4', + invocation_id='inv-4', + author='agent', + timestamp=12348, + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall(name='test_function') + ) + ] + ), + ), + ], +) + +MOCK_SESSION_WITH_EMPTY_EVENTS = Session( + app_name=MOCK_APP_NAME, + user_id=MOCK_USER_ID, + id=MOCK_SESSION_ID, + last_update_time=1000, +) + + +@pytest.fixture +def mock_openmemory_client(): + """Mock OpenMemory SDK client for testing.""" + with patch('google.adk.memory.open_memory_service.OpenMemory') as mock_client_constructor: + mock_client = MagicMock() + mock_client.add = MagicMock() + mock_client.query = MagicMock() + mock_client_constructor.return_value = mock_client + yield mock_client + + +@pytest.fixture +def memory_service(mock_openmemory_client): + """Create OpenMemoryService instance for testing.""" + return OpenMemoryService(base_url='http://localhost:3000', api_key='test-key') + + +@pytest.fixture +def memory_service_with_config(mock_openmemory_client): + """Create OpenMemoryService with custom config.""" + config = OpenMemoryServiceConfig( + search_top_k=5, + user_content_salience=0.9, + model_content_salience=0.6 + ) + return OpenMemoryService( + base_url='http://localhost:3000', + api_key='test-key', + config=config + ) + + +class TestOpenMemoryServiceConfig: + """Tests for OpenMemoryServiceConfig.""" + + def test_default_config(self): + """Test default configuration values.""" + config = OpenMemoryServiceConfig() + assert config.search_top_k == 10 + assert config.timeout == 30.0 + assert config.user_content_salience == 0.8 + assert config.model_content_salience == 0.7 + assert config.default_salience == 0.6 + assert config.enable_metadata_tags is True + + def test_custom_config(self): + """Test custom configuration values.""" + config = OpenMemoryServiceConfig( + search_top_k=20, + timeout=10.0, + user_content_salience=0.9, + model_content_salience=0.75, + default_salience=0.5, + enable_metadata_tags=False + ) + assert config.search_top_k == 20 + assert config.timeout == 10.0 + assert config.user_content_salience == 0.9 + assert config.model_content_salience == 0.75 + assert config.default_salience == 0.5 + assert config.enable_metadata_tags is False + + def test_config_validation_search_top_k(self): + """Test search_top_k validation.""" + with pytest.raises(Exception): # Pydantic validation error + OpenMemoryServiceConfig(search_top_k=0) + + with pytest.raises(Exception): + OpenMemoryServiceConfig(search_top_k=101) + + +class TestOpenMemoryService: + """Tests for OpenMemoryService.""" + + @pytest.mark.asyncio + async def test_add_session_to_memory_success(self, memory_service, mock_openmemory_client): + """Test successful addition of session memories.""" + mock_openmemory_client.add.return_value = {"id": "mem-1"} + + await memory_service.add_session_to_memory(MOCK_SESSION) + + # Should make 2 add calls (one per valid event) + assert mock_openmemory_client.add.call_count == 2 + + # Check first call (user event) + call_args = mock_openmemory_client.add.call_args_list[0] + request_data = call_args[0][0] + assert request_data['content'] == 'Hello, I like Python.' + assert 'session:session-1' in request_data['tags'] + assert request_data['metadata']['author'] == 'user' + assert request_data['salience'] == 0.8 # User content salience + + # Check second call (model event) + call_args = mock_openmemory_client.add.call_args_list[1] + request_data = call_args[0][0] + assert request_data['content'] == 'Python is a great programming language.' + assert request_data['metadata']['author'] == 'model' + assert request_data['salience'] == 0.7 # Model content salience + + @pytest.mark.asyncio + async def test_add_session_filters_empty_events( + self, memory_service, mock_openmemory_client + ): + """Test that events without content are filtered out.""" + await memory_service.add_session_to_memory(MOCK_SESSION_WITH_EMPTY_EVENTS) + + # Should make 0 add calls (no valid events) + assert mock_openmemory_client.add.call_count == 0 + + @pytest.mark.asyncio + async def test_add_session_uses_config_salience( + self, memory_service_with_config, mock_openmemory_client + ): + """Test that salience values from config are used.""" + mock_openmemory_client.add.return_value = {"id": "mem-1"} + + await memory_service_with_config.add_session_to_memory(MOCK_SESSION) + + # Check that custom salience values are used + call_args = mock_openmemory_client.add.call_args_list[0] + request_data = call_args[0][0] + assert request_data['salience'] == 0.9 # Custom user salience + + call_args = mock_openmemory_client.add.call_args_list[1] + request_data = call_args[0][0] + assert request_data['salience'] == 0.6 # Custom model salience + + @pytest.mark.asyncio + async def test_add_session_without_metadata_tags( + self, mock_openmemory_client + ): + """Test adding memories without metadata tags.""" + config = OpenMemoryServiceConfig(enable_metadata_tags=False) + memory_service = OpenMemoryService( + base_url='http://localhost:3000', config=config + ) + mock_openmemory_client.add.return_value = {"id": "mem-1"} + + await memory_service.add_session_to_memory(MOCK_SESSION) + + call_args = mock_openmemory_client.add.call_args_list[0] + request_data = call_args[0][0] + assert 'tags' not in request_data + + @pytest.mark.asyncio + async def test_add_session_error_handling(self, memory_service, mock_openmemory_client): + """Test error handling during memory addition.""" + mock_openmemory_client.add.side_effect = Exception('API Error') + + # Should not raise exception, just log error + await memory_service.add_session_to_memory(MOCK_SESSION) + + # Should still attempt to make add calls + assert mock_openmemory_client.add.call_count == 2 + + @pytest.mark.asyncio + async def test_search_memory_success(self, memory_service, mock_openmemory_client): + """Test successful memory search.""" + mock_openmemory_client.query.return_value = { + 'items': [ + { + 'content': 'Python is great', + 'metadata': { + 'author': 'user', + 'timestamp': 12345, + 'user_id': MOCK_USER_ID, + 'app_name': MOCK_APP_NAME, + } + }, + { + 'content': 'I like programming', + 'metadata': { + 'author': 'model', + 'timestamp': 12346, + 'user_id': MOCK_USER_ID, + 'app_name': MOCK_APP_NAME, + } + } + ] + } + + result = await memory_service.search_memory( + app_name=MOCK_APP_NAME, + user_id=MOCK_USER_ID, + query='Python programming' + ) + + # Verify API call + call_args = mock_openmemory_client.query.call_args + request_data = call_args[0][0] + assert request_data['query'] == 'Python programming' + assert request_data['top_k'] == 10 + assert 'tags' in request_data['filter'] + + # Verify results + assert len(result.memories) == 2 + assert result.memories[0].content.parts[0].text == 'Python is great' + assert result.memories[0].author == 'user' + assert result.memories[1].content.parts[0].text == 'I like programming' + assert result.memories[1].author == 'model' + + @pytest.mark.asyncio + async def test_search_memory_applies_filters( + self, memory_service, mock_openmemory_client + ): + """Test that app_name/user_id filters are applied.""" + mock_openmemory_client.query.return_value = { + 'items': [ + { + 'content': 'Python is great', + 'metadata': { + 'author': 'user', + 'timestamp': 12345, + 'user_id': 'different-user', # Different user + 'app_name': MOCK_APP_NAME, + } + }, + { + 'content': 'I like programming', + 'metadata': { + 'author': 'model', + 'timestamp': 12346, + 'user_id': MOCK_USER_ID, # Correct user + 'app_name': MOCK_APP_NAME, + } + } + ] + } + + result = await memory_service.search_memory( + app_name=MOCK_APP_NAME, + user_id=MOCK_USER_ID, + query='test query' + ) + + # Should filter out the different user's result + assert len(result.memories) == 1 + assert result.memories[0].content.parts[0].text == 'I like programming' + + @pytest.mark.asyncio + async def test_search_memory_respects_top_k( + self, memory_service_with_config, mock_openmemory_client + ): + """Test that config.search_top_k is used.""" + mock_openmemory_client.query.return_value = {'items': []} + + await memory_service_with_config.search_memory( + app_name=MOCK_APP_NAME, + user_id=MOCK_USER_ID, + query='test query' + ) + + call_args = mock_openmemory_client.query.call_args + request_data = call_args[0][0] + assert request_data['top_k'] == 5 # Custom config value + + @pytest.mark.asyncio + async def test_search_memory_error_handling( + self, memory_service, mock_openmemory_client + ): + """Test graceful error handling during memory search.""" + mock_openmemory_client.query.side_effect = Exception('API Error') + + result = await memory_service.search_memory( + app_name=MOCK_APP_NAME, + user_id=MOCK_USER_ID, + query='test query' + ) + + # Should return empty results on error + assert len(result.memories) == 0 + + @pytest.mark.asyncio + async def test_close_client(self, memory_service): + """Test client cleanup (SDK handles automatically).""" + # SDK handles cleanup automatically, so close() is a no-op + await memory_service.close() + # Should complete without error + From da9fa068ad3d4b0104c1424dc0d0613b969c0dc7 Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 12:25:34 +0530 Subject: [PATCH 2/9] fix: Update timestamp handling and streamline memory addition process in OpenMemoryService - Changed the timestamp assignment to use event.timestamp instead of a formatted string. - Refactored the memory addition logic to ensure user_id is included as a top-level field in the payload for better server-side filtering. - Improved error handling and logging during memory addition for clarity and debugging purposes. These changes enhance the accuracy of memory data and improve the overall efficiency of the memory service operations. --- src/google/adk/memory/open_memory_service.py | 43 ++++++++++---------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/src/google/adk/memory/open_memory_service.py b/src/google/adk/memory/open_memory_service.py index fe78356de9..dc78d20a8a 100644 --- a/src/google/adk/memory/open_memory_service.py +++ b/src/google/adk/memory/open_memory_service.py @@ -152,10 +152,9 @@ def _prepare_memory_data( "event_id": event.id, "invocation_id": event.invocation_id, "author": event.author, - "timestamp": timestamp_str, + "timestamp": event.timestamp, "source": "adk_session" } - metadata = {k: v for k, v in metadata.items() if v is not None} memory_data = { "content": enriched_content, @@ -186,19 +185,21 @@ async def add_session_to_memory(self, session: Session): """ memories_added = 0 - headers = {"Content-Type": "application/json"} - if self._api_key: - headers["Authorization"] = f"Bearer {self._api_key}" - - async with httpx.AsyncClient(timeout=self._config.timeout) as http_client: - for event in session.events: - content_text = _extract_text_from_event(event) - if not content_text: - continue - - memory_data = self._prepare_memory_data(event, content_text, session) - - try: + for event in session.events: + content_text = _extract_text_from_event(event) + if not content_text: + continue + + memory_data = self._prepare_memory_data(event, content_text, session) + + try: + # Use direct HTTP to pass user_id as top-level field (not just in metadata) + # This ensures server-side filtering works correctly + async with httpx.AsyncClient(timeout=self._config.timeout) as http_client: + headers = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + # Include user_id as separate field for database storage and filtering payload = { "content": memory_data["content"], @@ -207,18 +208,18 @@ async def add_session_to_memory(self, session: Session): "salience": memory_data.get("salience", 0.5), "user_id": session.user_id # Separate field for DB column } - + response = await http_client.post( f"{self._base_url}/memory/add", json=payload, headers=headers ) response.raise_for_status() - - memories_added += 1 - logger.debug("Added memory for event %s", event.id) - except Exception as e: - logger.error("Failed to add memory for event %s: %s", event.id, e) + + memories_added += 1 + logger.debug("Added memory for event %s", event.id) + except Exception as e: + logger.error("Failed to add memory for event %s: %s", event.id, e) logger.info( "Added %d memories from session %s", memories_added, session.id From 72961db374feb1ec5106e0b2a889a2b38f5b37e6 Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 14:49:14 +0530 Subject: [PATCH 3/9] refactor: Improve memory service performance and update agent instructions - Refactored the memory addition process in OpenMemoryService to create the HTTP client once for all events, enhancing performance. - Updated the agent's instruction format for better readability and maintainability. - Changed import statements in the README to reflect the updated Runner class. These changes optimize the memory service operations and improve the clarity of the agent's instructions. --- contributing/samples/open_memory/README.md | 4 +- contributing/samples/open_memory/agent.py | 8 +- src/google/adk/memory/open_memory_service.py | 40 ++-- .../memory/test_open_memory_service.py | 172 ++++++++---------- 4 files changed, 105 insertions(+), 119 deletions(-) diff --git a/contributing/samples/open_memory/README.md b/contributing/samples/open_memory/README.md index 0abaac8267..4a72c6eab3 100644 --- a/contributing/samples/open_memory/README.md +++ b/contributing/samples/open_memory/README.md @@ -64,7 +64,7 @@ OPENMEMORY_BASE_URL=http://localhost:3000 ```python from google.adk.memory import OpenMemoryService -from google.adk.runners import InMemoryRunner +from google.adk.runners import Runner # Create OpenMemory service with defaults memory_service = OpenMemoryService( @@ -72,7 +72,7 @@ memory_service = OpenMemoryService( ) # Use with runner -runner = InMemoryRunner( +runner = Runner( app_name="my_app", agent=my_agent, memory_service=memory_service diff --git a/contributing/samples/open_memory/agent.py b/contributing/samples/open_memory/agent.py index 8a61a7db68..b5678184de 100644 --- a/contributing/samples/open_memory/agent.py +++ b/contributing/samples/open_memory/agent.py @@ -30,10 +30,10 @@ def update_current_time(callback_context: CallbackContext): name='open_memory_agent', description='agent that has access to memory tools with OpenMemory.', before_agent_callback=update_current_time, - instruction="""\ - You are an agent that help user answer questions. You have access to memory tools. - You can use the memory tools to look up the information in the memory. Current time: {_time} - """, + instruction=( + 'You are an agent that help user answer questions. You have access to memory tools.\n' + 'You can use the memory tools to look up the information in the memory. Current time: {_time}' + ), tools=[load_memory_tool, preload_memory_tool], ) diff --git a/src/google/adk/memory/open_memory_service.py b/src/google/adk/memory/open_memory_service.py index dc78d20a8a..9a942aaee4 100644 --- a/src/google/adk/memory/open_memory_service.py +++ b/src/google/adk/memory/open_memory_service.py @@ -185,20 +185,22 @@ async def add_session_to_memory(self, session: Session): """ memories_added = 0 - for event in session.events: - content_text = _extract_text_from_event(event) - if not content_text: - continue - - memory_data = self._prepare_memory_data(event, content_text, session) - - try: - # Use direct HTTP to pass user_id as top-level field (not just in metadata) - # This ensures server-side filtering works correctly - async with httpx.AsyncClient(timeout=self._config.timeout) as http_client: - headers = {"Content-Type": "application/json"} - if self._api_key: - headers["Authorization"] = f"Bearer {self._api_key}" + # Create HTTP client once for all events to improve performance + async with httpx.AsyncClient(timeout=self._config.timeout) as http_client: + headers = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + + for event in session.events: + content_text = _extract_text_from_event(event) + if not content_text: + continue + + memory_data = self._prepare_memory_data(event, content_text, session) + + try: + # Use direct HTTP to pass user_id as top-level field + # This ensures server-side filtering works correctly # Include user_id as separate field for database storage and filtering payload = { @@ -215,11 +217,11 @@ async def add_session_to_memory(self, session: Session): headers=headers ) response.raise_for_status() - - memories_added += 1 - logger.debug("Added memory for event %s", event.id) - except Exception as e: - logger.error("Failed to add memory for event %s: %s", event.id, e) + + memories_added += 1 + logger.debug("Added memory for event %s", event.id) + except Exception as e: + logger.error("Failed to add memory for event %s: %s", event.id, e) logger.info( "Added %d memories from session %s", memories_added, session.id diff --git a/tests/unittests/memory/test_open_memory_service.py b/tests/unittests/memory/test_open_memory_service.py index 88ef1cff39..d171fc410d 100644 --- a/tests/unittests/memory/test_open_memory_service.py +++ b/tests/unittests/memory/test_open_memory_service.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock from unittest.mock import patch from google.adk.events.event import Event @@ -83,24 +83,28 @@ @pytest.fixture -def mock_openmemory_client(): - """Mock OpenMemory SDK client for testing.""" - with patch('google.adk.memory.open_memory_service.OpenMemory') as mock_client_constructor: +def mock_httpx_client(): + """Mock httpx.AsyncClient for testing.""" + with patch('google.adk.memory.open_memory_service.httpx.AsyncClient') as mock_client_class: mock_client = MagicMock() - mock_client.add = MagicMock() - mock_client.query = MagicMock() - mock_client_constructor.return_value = mock_client + mock_response = MagicMock() + mock_response.json.return_value = {"matches": []} + mock_response.raise_for_status = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + mock_client_class.return_value = mock_client yield mock_client @pytest.fixture -def memory_service(mock_openmemory_client): +def memory_service(mock_httpx_client): """Create OpenMemoryService instance for testing.""" return OpenMemoryService(base_url='http://localhost:3000', api_key='test-key') @pytest.fixture -def memory_service_with_config(mock_openmemory_client): +def memory_service_with_config(mock_httpx_client): """Create OpenMemoryService with custom config.""" config = OpenMemoryServiceConfig( search_top_k=5, @@ -157,111 +161,100 @@ class TestOpenMemoryService: """Tests for OpenMemoryService.""" @pytest.mark.asyncio - async def test_add_session_to_memory_success(self, memory_service, mock_openmemory_client): + async def test_add_session_to_memory_success(self, memory_service, mock_httpx_client): """Test successful addition of session memories.""" - mock_openmemory_client.add.return_value = {"id": "mem-1"} - await memory_service.add_session_to_memory(MOCK_SESSION) - # Should make 2 add calls (one per valid event) - assert mock_openmemory_client.add.call_count == 2 + # Should make 2 POST calls (one per valid event) + assert mock_httpx_client.post.call_count == 2 # Check first call (user event) - call_args = mock_openmemory_client.add.call_args_list[0] - request_data = call_args[0][0] - assert request_data['content'] == 'Hello, I like Python.' + call_args = mock_httpx_client.post.call_args_list[0] + request_data = call_args.kwargs['json'] + assert '[Author: user' in request_data['content'] + assert 'Hello, I like Python.' in request_data['content'] assert 'session:session-1' in request_data['tags'] assert request_data['metadata']['author'] == 'user' assert request_data['salience'] == 0.8 # User content salience # Check second call (model event) - call_args = mock_openmemory_client.add.call_args_list[1] - request_data = call_args[0][0] - assert request_data['content'] == 'Python is a great programming language.' + call_args = mock_httpx_client.post.call_args_list[1] + request_data = call_args.kwargs['json'] + assert '[Author: model' in request_data['content'] + assert 'Python is a great programming language.' in request_data['content'] assert request_data['metadata']['author'] == 'model' assert request_data['salience'] == 0.7 # Model content salience @pytest.mark.asyncio async def test_add_session_filters_empty_events( - self, memory_service, mock_openmemory_client + self, memory_service, mock_httpx_client ): """Test that events without content are filtered out.""" await memory_service.add_session_to_memory(MOCK_SESSION_WITH_EMPTY_EVENTS) - # Should make 0 add calls (no valid events) - assert mock_openmemory_client.add.call_count == 0 + # Should make 0 POST calls (no valid events) + assert mock_httpx_client.post.call_count == 0 @pytest.mark.asyncio async def test_add_session_uses_config_salience( - self, memory_service_with_config, mock_openmemory_client + self, memory_service_with_config, mock_httpx_client ): """Test that salience values from config are used.""" - mock_openmemory_client.add.return_value = {"id": "mem-1"} - await memory_service_with_config.add_session_to_memory(MOCK_SESSION) # Check that custom salience values are used - call_args = mock_openmemory_client.add.call_args_list[0] - request_data = call_args[0][0] + call_args = mock_httpx_client.post.call_args_list[0] + request_data = call_args.kwargs['json'] assert request_data['salience'] == 0.9 # Custom user salience - call_args = mock_openmemory_client.add.call_args_list[1] - request_data = call_args[0][0] + call_args = mock_httpx_client.post.call_args_list[1] + request_data = call_args.kwargs['json'] assert request_data['salience'] == 0.6 # Custom model salience @pytest.mark.asyncio async def test_add_session_without_metadata_tags( - self, mock_openmemory_client + self, mock_httpx_client ): """Test adding memories without metadata tags.""" config = OpenMemoryServiceConfig(enable_metadata_tags=False) memory_service = OpenMemoryService( base_url='http://localhost:3000', config=config ) - mock_openmemory_client.add.return_value = {"id": "mem-1"} await memory_service.add_session_to_memory(MOCK_SESSION) - call_args = mock_openmemory_client.add.call_args_list[0] - request_data = call_args[0][0] - assert 'tags' not in request_data + call_args = mock_httpx_client.post.call_args_list[0] + request_data = call_args.kwargs['json'] + assert request_data.get('tags', []) == [] @pytest.mark.asyncio - async def test_add_session_error_handling(self, memory_service, mock_openmemory_client): + async def test_add_session_error_handling(self, memory_service, mock_httpx_client): """Test error handling during memory addition.""" - mock_openmemory_client.add.side_effect = Exception('API Error') + mock_httpx_client.post.side_effect = Exception('API Error') # Should not raise exception, just log error await memory_service.add_session_to_memory(MOCK_SESSION) - # Should still attempt to make add calls - assert mock_openmemory_client.add.call_count == 2 + # Should still attempt to make POST calls + assert mock_httpx_client.post.call_count == 2 @pytest.mark.asyncio - async def test_search_memory_success(self, memory_service, mock_openmemory_client): + async def test_search_memory_success(self, memory_service, mock_httpx_client): """Test successful memory search.""" - mock_openmemory_client.query.return_value = { - 'items': [ + # Mock response with enriched content format + mock_response = MagicMock() + mock_response.json.return_value = { + 'matches': [ { - 'content': 'Python is great', - 'metadata': { - 'author': 'user', - 'timestamp': 12345, - 'user_id': MOCK_USER_ID, - 'app_name': MOCK_APP_NAME, - } + 'content': '[Author: user, Time: 2025-01-01T00:00:00] Python is great', }, { - 'content': 'I like programming', - 'metadata': { - 'author': 'model', - 'timestamp': 12346, - 'user_id': MOCK_USER_ID, - 'app_name': MOCK_APP_NAME, - } + 'content': '[Author: model, Time: 2025-01-01T00:01:00] I like programming', } ] } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.post = AsyncMock(return_value=mock_response) result = await memory_service.search_memory( app_name=MOCK_APP_NAME, @@ -270,13 +263,14 @@ async def test_search_memory_success(self, memory_service, mock_openmemory_clien ) # Verify API call - call_args = mock_openmemory_client.query.call_args - request_data = call_args[0][0] + call_args = mock_httpx_client.post.call_args + request_data = call_args.kwargs['json'] assert request_data['query'] == 'Python programming' assert request_data['top_k'] == 10 - assert 'tags' in request_data['filter'] + assert request_data['filters']['user_id'] == MOCK_USER_ID + assert request_data['filters']['app_name'] == MOCK_APP_NAME - # Verify results + # Verify results (content should be cleaned of metadata prefix) assert len(result.memories) == 2 assert result.memories[0].content.parts[0].text == 'Python is great' assert result.memories[0].author == 'user' @@ -285,31 +279,20 @@ async def test_search_memory_success(self, memory_service, mock_openmemory_clien @pytest.mark.asyncio async def test_search_memory_applies_filters( - self, memory_service, mock_openmemory_client + self, memory_service, mock_httpx_client ): """Test that app_name/user_id filters are applied.""" - mock_openmemory_client.query.return_value = { - 'items': [ - { - 'content': 'Python is great', - 'metadata': { - 'author': 'user', - 'timestamp': 12345, - 'user_id': 'different-user', # Different user - 'app_name': MOCK_APP_NAME, - } - }, + # Mock response - server-side filtering ensures only matching results + mock_response = MagicMock() + mock_response.json.return_value = { + 'matches': [ { - 'content': 'I like programming', - 'metadata': { - 'author': 'model', - 'timestamp': 12346, - 'user_id': MOCK_USER_ID, # Correct user - 'app_name': MOCK_APP_NAME, - } + 'content': '[Author: model, Time: 2025-01-01T00:01:00] I like programming', } ] } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.post = AsyncMock(return_value=mock_response) result = await memory_service.search_memory( app_name=MOCK_APP_NAME, @@ -317,16 +300,25 @@ async def test_search_memory_applies_filters( query='test query' ) - # Should filter out the different user's result + # Verify filters were passed correctly + call_args = mock_httpx_client.post.call_args + request_data = call_args.kwargs['json'] + assert request_data['filters']['user_id'] == MOCK_USER_ID + assert request_data['filters']['app_name'] == MOCK_APP_NAME + + # Should return filtered results assert len(result.memories) == 1 assert result.memories[0].content.parts[0].text == 'I like programming' @pytest.mark.asyncio async def test_search_memory_respects_top_k( - self, memory_service_with_config, mock_openmemory_client + self, memory_service_with_config, mock_httpx_client ): """Test that config.search_top_k is used.""" - mock_openmemory_client.query.return_value = {'items': []} + mock_response = MagicMock() + mock_response.json.return_value = {'matches': []} + mock_response.raise_for_status = MagicMock() + mock_httpx_client.post = AsyncMock(return_value=mock_response) await memory_service_with_config.search_memory( app_name=MOCK_APP_NAME, @@ -334,16 +326,16 @@ async def test_search_memory_respects_top_k( query='test query' ) - call_args = mock_openmemory_client.query.call_args - request_data = call_args[0][0] + call_args = mock_httpx_client.post.call_args + request_data = call_args.kwargs['json'] assert request_data['top_k'] == 5 # Custom config value @pytest.mark.asyncio async def test_search_memory_error_handling( - self, memory_service, mock_openmemory_client + self, memory_service, mock_httpx_client ): """Test graceful error handling during memory search.""" - mock_openmemory_client.query.side_effect = Exception('API Error') + mock_httpx_client.post.side_effect = Exception('API Error') result = await memory_service.search_memory( app_name=MOCK_APP_NAME, @@ -353,11 +345,3 @@ async def test_search_memory_error_handling( # Should return empty results on error assert len(result.memories) == 0 - - @pytest.mark.asyncio - async def test_close_client(self, memory_service): - """Test client cleanup (SDK handles automatically).""" - # SDK handles cleanup automatically, so close() is a no-op - await memory_service.close() - # Should complete without error - From 4c14e77a1f88bb8da6ad05bf45643ddd3117e19c Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 15:31:50 +0530 Subject: [PATCH 4/9] feat: Add OpenMemory service connection options and factory method - Updated the agent's instructions to include connection details for OpenMemory, allowing users to connect using various URI formats. - Implemented an `openmemory_factory` method to create OpenMemoryService instances from specified URIs, enhancing flexibility in service registration. - Added a no-op `close` method in OpenMemoryService for API consistency. These changes improve the usability and configurability of the OpenMemory service within the application. --- src/google/adk/cli/cli_tools_click.py | 5 +- src/google/adk/cli/service_registry.py | 63 ++++++++++++++++++++ src/google/adk/memory/open_memory_service.py | 9 +++ 3 files changed, 76 insertions(+), 1 deletion(-) diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 06f567a037..b3be833034 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -898,7 +898,10 @@ def decorator(func): - Use 'agentengine://' to connect to Agent Engine sessions. can either be the full qualified resource name 'projects/abc/locations/us-central1/reasoningEngines/123' or - the resource id '123'."""), + the resource id '123'. + - Use 'openmemory://:' to connect to OpenMemory. + Example: 'openmemory://localhost:3000' or + 'openmemory://localhost:3000?api_key=secret'."""), default=None, ) @functools.wraps(func) diff --git a/src/google/adk/cli/service_registry.py b/src/google/adk/cli/service_registry.py index bc95bad285..0e4cbc0844 100644 --- a/src/google/adk/cli/service_registry.py +++ b/src/google/adk/cli/service_registry.py @@ -18,6 +18,7 @@ from typing import Any from typing import Dict from typing import Protocol +from urllib.parse import parse_qs from urllib.parse import urlparse from ..artifacts.base_artifact_service import BaseArtifactService @@ -210,8 +211,70 @@ def agentengine_memory_factory(uri: str, **kwargs): ) return VertexAiMemoryBankService(**params) + def openmemory_factory(uri: str, **kwargs): + """Factory for OpenMemory service from URI. + + Supported URI formats: + - openmemory://localhost:3000 -> http://localhost:3000 + - openmemory://https://example.com -> https://example.com + - openmemory://localhost:3000?api_key=secret -> http://localhost:3000 with API key + + Args: + uri: URI in format openmemory://: or openmemory:// + **kwargs: Additional arguments (agents_dir is ignored) + + Returns: + OpenMemoryService instance + """ + try: + from ..memory.open_memory_service import OpenMemoryService + except ImportError: + raise ImportError( + "OpenMemoryService requires httpx. Install it with:" + " pip install google-adk[openmemory]" + ) + + parsed = urlparse(uri) + + # Extract base URL + # The netloc contains the host:port, and path may contain additional path + # If path looks like a full URL (starts with //), use it + # Otherwise, construct http:// + netloc = parsed.netloc or "" + path = parsed.path + + # Check if path contains a full URL (e.g., openmemory://http://localhost:3000) + if path.startswith("//"): + # Extract the URL from path (e.g., //http://localhost:3000 -> http://localhost:3000) + full_url = path.lstrip("/") + if full_url.startswith(("http://", "https://")): + base_url = full_url + else: + base_url = f"http://{full_url}" + elif netloc.startswith(("http://", "https://")): + # Netloc itself is a full URL (shouldn't happen with proper URL parsing, but handle it) + base_url = netloc + else: + # Construct URL from netloc and path + if netloc: + base_url = f"http://{netloc}{path}" + else: + raise ValueError( + f"Invalid OpenMemory URI: {uri}. Expected format:" + " openmemory://localhost:3000 or openmemory://http://localhost:3000" + ) + + # Extract API key from query parameters if present + api_key = None + if parsed.query: + query_params = parse_qs(parsed.query) + api_key = query_params.get("api_key", [None])[0] + + return OpenMemoryService(base_url=base_url, api_key=api_key) + registry.register_memory_service("rag", rag_memory_factory) registry.register_memory_service("agentengine", agentengine_memory_factory) + registry.register_memory_service("openmemory", openmemory_factory) # Global registry instance diff --git a/src/google/adk/memory/open_memory_service.py b/src/google/adk/memory/open_memory_service.py index 9a942aaee4..5010e97fae 100644 --- a/src/google/adk/memory/open_memory_service.py +++ b/src/google/adk/memory/open_memory_service.py @@ -359,6 +359,15 @@ async def search_memory( logger.error("Failed to search memories: %s", e) return SearchMemoryResponse(memories=[]) + async def close(self): + """Close the memory service and cleanup resources. + + This method is provided for API consistency. Since httpx.AsyncClient + is used as a context manager in all operations, cleanup is handled + automatically. This method is a no-op and can be safely called or omitted. + """ + pass + class OpenMemoryServiceConfig(BaseModel): """Configuration for OpenMemory service behavior. From ebb7f88abcf2f88238310bda3f130ab36667c4a9 Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 15:42:26 +0530 Subject: [PATCH 5/9] refactor: Enhance URI handling in service registration - Improved the URI parsing logic in the _register_builtin_services function to support various OpenMemory URI formats, including handling cases where the netloc is a scheme. - Updated comments for clarity on how different URI formats are processed, ensuring better understanding and maintainability of the code. These changes enhance the flexibility and robustness of service registration for OpenMemory connections. --- src/google/adk/cli/service_registry.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/google/adk/cli/service_registry.py b/src/google/adk/cli/service_registry.py index 0e4cbc0844..f0aa34a705 100644 --- a/src/google/adk/cli/service_registry.py +++ b/src/google/adk/cli/service_registry.py @@ -237,15 +237,22 @@ def openmemory_factory(uri: str, **kwargs): parsed = urlparse(uri) # Extract base URL - # The netloc contains the host:port, and path may contain additional path - # If path looks like a full URL (starts with //), use it - # Otherwise, construct http:// + # Handle different URI formats: + # - openmemory://localhost:3000 -> http://localhost:3000 + # - openmemory://https://example.com -> https://example.com + # - openmemory://http://localhost:3000 -> http://localhost:3000 netloc = parsed.netloc or "" path = parsed.path - # Check if path contains a full URL (e.g., openmemory://http://localhost:3000) - if path.startswith("//"): - # Extract the URL from path (e.g., //http://localhost:3000 -> http://localhost:3000) + # Check if netloc is a scheme (e.g., "https:" or "http:") + # This happens when URI is like openmemory://https://example.com + if netloc.endswith(":") and path.startswith("//"): + # Reconstruct the full URL: scheme from netloc + path + scheme = netloc.rstrip(":") + # path is like "//example.com", we want "https://example.com" + base_url = f"{scheme}://{path[2:]}" # Remove "//" prefix + elif path.startswith("//"): + # Path contains a full URL (e.g., openmemory:////http://localhost:3000) full_url = path.lstrip("/") if full_url.startswith(("http://", "https://")): base_url = full_url @@ -255,13 +262,13 @@ def openmemory_factory(uri: str, **kwargs): # Netloc itself is a full URL (shouldn't happen with proper URL parsing, but handle it) base_url = netloc else: - # Construct URL from netloc and path + # Construct URL from netloc and path (default case) if netloc: base_url = f"http://{netloc}{path}" else: raise ValueError( f"Invalid OpenMemory URI: {uri}. Expected format:" - " openmemory://localhost:3000 or openmemory://http://localhost:3000" + " openmemory://localhost:3000 or openmemory://https://example.com" ) # Extract API key from query parameters if present From ce8f628fbf6d99b51b15f2e419a91843c54c9862 Mon Sep 17 00:00:00 2001 From: Rakshith G <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 16:03:56 +0530 Subject: [PATCH 6/9] Update tests/unittests/memory/test_open_memory_service.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/unittests/memory/test_open_memory_service.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unittests/memory/test_open_memory_service.py b/tests/unittests/memory/test_open_memory_service.py index d171fc410d..d93781783f 100644 --- a/tests/unittests/memory/test_open_memory_service.py +++ b/tests/unittests/memory/test_open_memory_service.py @@ -266,9 +266,9 @@ async def test_search_memory_success(self, memory_service, mock_httpx_client): call_args = mock_httpx_client.post.call_args request_data = call_args.kwargs['json'] assert request_data['query'] == 'Python programming' - assert request_data['top_k'] == 10 - assert request_data['filters']['user_id'] == MOCK_USER_ID - assert request_data['filters']['app_name'] == MOCK_APP_NAME + assert request_data['k'] == 10 + assert request_data['filter']['user_id'] == MOCK_USER_ID + assert request_data['filter']['tags'] == [f'app:{MOCK_APP_NAME}'] # Verify results (content should be cleaned of metadata prefix) assert len(result.memories) == 2 From 5392512b3513b3fdd73b4d0ca199f10545c8a384 Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 19:10:08 +0530 Subject: [PATCH 7/9] Address review comments: simplify URI parsing and tag creation logic - Simplified openmemory URI parsing logic to be more maintainable - Improved tag creation to avoid empty author tags - Fixed test assertions to match actual API field names (k vs top_k, filter vs filters) --- src/google/adk/cli/service_registry.py | 45 ++++++------------- src/google/adk/memory/open_memory_service.py | 8 ++-- .../memory/test_open_memory_service.py | 8 ++-- 3 files changed, 22 insertions(+), 39 deletions(-) diff --git a/src/google/adk/cli/service_registry.py b/src/google/adk/cli/service_registry.py index f0aa34a705..8785217fbe 100644 --- a/src/google/adk/cli/service_registry.py +++ b/src/google/adk/cli/service_registry.py @@ -237,39 +237,22 @@ def openmemory_factory(uri: str, **kwargs): parsed = urlparse(uri) # Extract base URL - # Handle different URI formats: - # - openmemory://localhost:3000 -> http://localhost:3000 - # - openmemory://https://example.com -> https://example.com - # - openmemory://http://localhost:3000 -> http://localhost:3000 - netloc = parsed.netloc or "" - path = parsed.path + # The part after "openmemory://" + location_part = uri[len("openmemory://"):] - # Check if netloc is a scheme (e.g., "https:" or "http:") - # This happens when URI is like openmemory://https://example.com - if netloc.endswith(":") and path.startswith("//"): - # Reconstruct the full URL: scheme from netloc + path - scheme = netloc.rstrip(":") - # path is like "//example.com", we want "https://example.com" - base_url = f"{scheme}://{path[2:]}" # Remove "//" prefix - elif path.startswith("//"): - # Path contains a full URL (e.g., openmemory:////http://localhost:3000) - full_url = path.lstrip("/") - if full_url.startswith(("http://", "https://")): - base_url = full_url - else: - base_url = f"http://{full_url}" - elif netloc.startswith(("http://", "https://")): - # Netloc itself is a full URL (shouldn't happen with proper URL parsing, but handle it) - base_url = netloc + # Remove query string for base_url construction + base_url_part = location_part.split('?')[0].rstrip('/') + + if not base_url_part: + raise ValueError( + f"Invalid OpenMemory URI: {uri}. Expected format:" + " openmemory://localhost:3000 or openmemory://https://example.com" + ) + + if not base_url_part.startswith(("http://", "https://")): + base_url = f"http://{base_url_part}" else: - # Construct URL from netloc and path (default case) - if netloc: - base_url = f"http://{netloc}{path}" - else: - raise ValueError( - f"Invalid OpenMemory URI: {uri}. Expected format:" - " openmemory://localhost:3000 or openmemory://https://example.com" - ) + base_url = base_url_part # Extract API key from query parameters if present api_key = None diff --git a/src/google/adk/memory/open_memory_service.py b/src/google/adk/memory/open_memory_service.py index 5010e97fae..0754107ddf 100644 --- a/src/google/adk/memory/open_memory_service.py +++ b/src/google/adk/memory/open_memory_service.py @@ -163,13 +163,13 @@ def _prepare_memory_data( } if self._config.enable_metadata_tags: - memory_data["tags"] = [ + tags = [ f"session:{session.id}", f"app:{session.app_name}", - f"author:{event.author}" if event.author else None ] - # Remove None values - memory_data["tags"] = [t for t in memory_data["tags"] if t] + if event.author: + tags.append(f"author:{event.author}") + memory_data["tags"] = tags return memory_data diff --git a/tests/unittests/memory/test_open_memory_service.py b/tests/unittests/memory/test_open_memory_service.py index d93781783f..5c28a2bc5c 100644 --- a/tests/unittests/memory/test_open_memory_service.py +++ b/tests/unittests/memory/test_open_memory_service.py @@ -268,7 +268,7 @@ async def test_search_memory_success(self, memory_service, mock_httpx_client): assert request_data['query'] == 'Python programming' assert request_data['k'] == 10 assert request_data['filter']['user_id'] == MOCK_USER_ID - assert request_data['filter']['tags'] == [f'app:{MOCK_APP_NAME}'] + assert f"app:{MOCK_APP_NAME}" in request_data['filter']['tags'] # Verify results (content should be cleaned of metadata prefix) assert len(result.memories) == 2 @@ -303,8 +303,8 @@ async def test_search_memory_applies_filters( # Verify filters were passed correctly call_args = mock_httpx_client.post.call_args request_data = call_args.kwargs['json'] - assert request_data['filters']['user_id'] == MOCK_USER_ID - assert request_data['filters']['app_name'] == MOCK_APP_NAME + assert request_data['filter']['user_id'] == MOCK_USER_ID + assert f"app:{MOCK_APP_NAME}" in request_data['filter']['tags'] # Should return filtered results assert len(result.memories) == 1 @@ -328,7 +328,7 @@ async def test_search_memory_respects_top_k( call_args = mock_httpx_client.post.call_args request_data = call_args.kwargs['json'] - assert request_data['top_k'] == 5 # Custom config value + assert request_data['k'] == 5 # Custom config value @pytest.mark.asyncio async def test_search_memory_error_handling( From c598171879ae744bbbaf5cb77a7d6a88985492b9 Mon Sep 17 00:00:00 2001 From: Rakshith G <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 19:17:27 +0530 Subject: [PATCH 8/9] Update contributing/samples/open_memory/agent.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- contributing/samples/open_memory/agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contributing/samples/open_memory/agent.py b/contributing/samples/open_memory/agent.py index b5678184de..1e420b548e 100644 --- a/contributing/samples/open_memory/agent.py +++ b/contributing/samples/open_memory/agent.py @@ -31,7 +31,7 @@ def update_current_time(callback_context: CallbackContext): description='agent that has access to memory tools with OpenMemory.', before_agent_callback=update_current_time, instruction=( - 'You are an agent that help user answer questions. You have access to memory tools.\n' + 'You are an agent that helps user answer questions. You have access to memory tools.\n' 'You can use the memory tools to look up the information in the memory. Current time: {_time}' ), tools=[load_memory_tool, preload_memory_tool], From 1d088a32c77d29ac27e33e127f6ba1fcde692c3d Mon Sep 17 00:00:00 2001 From: rakshith-git <62304358+rakshith-git@users.noreply.github.com> Date: Tue, 4 Nov 2025 19:18:30 +0530 Subject: [PATCH 9/9] Improve error handling with specific httpx exceptions - Replace generic Exception catches with specific httpx.HTTPStatusError and httpx.RequestError handlers - Provide more detailed error messages with status codes and response text - Keep generic Exception handler as fallback for unexpected errors --- src/google/adk/memory/open_memory_service.py | 25 ++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/src/google/adk/memory/open_memory_service.py b/src/google/adk/memory/open_memory_service.py index 0754107ddf..6f3d843728 100644 --- a/src/google/adk/memory/open_memory_service.py +++ b/src/google/adk/memory/open_memory_service.py @@ -220,8 +220,19 @@ async def add_session_to_memory(self, session: Session): memories_added += 1 logger.debug("Added memory for event %s", event.id) + except httpx.HTTPStatusError as e: + logger.error( + "Failed to add memory for event %s due to HTTP error: %s - %s", + event.id, + e.response.status_code, + e.response.text, + ) + except httpx.RequestError as e: + logger.error( + "Failed to add memory for event %s due to request error: %s", event.id, e + ) except Exception as e: - logger.error("Failed to add memory for event %s: %s", event.id, e) + logger.error("Failed to add memory for event %s due to unexpected error: %s", event.id, e) logger.info( "Added %d memories from session %s", memories_added, session.id @@ -355,8 +366,18 @@ async def search_memory( logger.info("Found %d memories for query: '%s'", len(memories), query) return SearchMemoryResponse(memories=memories) + except httpx.HTTPStatusError as e: + logger.error( + "Failed to search memories due to HTTP error: %s - %s", + e.response.status_code, + e.response.text, + ) + return SearchMemoryResponse(memories=[]) + except httpx.RequestError as e: + logger.error("Failed to search memories due to request error: %s", e) + return SearchMemoryResponse(memories=[]) except Exception as e: - logger.error("Failed to search memories: %s", e) + logger.error("Failed to search memories due to unexpected error: %s", e) return SearchMemoryResponse(memories=[]) async def close(self):