google-adk-extras 0.1.1__py3-none-any.whl

google_adk_extras/memory/yaml_file_memory_service.py

@@ -0,0 +1,176 @@
+"""YAML file-based memory service implementation."""
+
+import os
+import json
+import logging
+from typing import Optional, List
+from pathlib import Path
+import re
+from datetime import datetime
+
+import yaml
+
+from google.genai import types
+from .base_custom_memory_service import BaseCustomMemoryService
+
+
+logger = logging.getLogger('google_adk_extras.' + __name__)
+
+
+class YamlFileMemoryService(BaseCustomMemoryService):
+    """YAML file-based memory service implementation."""
+
+    def __init__(self, base_directory: str = "./memory"):
+        """Initialize the YAML file memory service.
+        
+        Args:
+            base_directory: Base directory for storing memory files
+        """
+        super().__init__()
+        self.base_directory = Path(base_directory)
+        # Create base directory if it doesn't exist
+        self.base_directory.mkdir(parents=True, exist_ok=True)
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the file system memory service."""
+        # Ensure base directory exists
+        self.base_directory.mkdir(parents=True, exist_ok=True)
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up resources (no cleanup needed for file-based service)."""
+        pass
+
+    def _get_memory_directory(self, app_name: str, user_id: str) -> Path:
+        """Generate directory path for memory entries."""
+        directory = self.base_directory / app_name / user_id
+        directory.mkdir(parents=True, exist_ok=True)
+        return directory
+
+    def _get_memory_file_path(self, app_name: str, user_id: str, memory_id: str) -> Path:
+        """Generate file path for a memory entry."""
+        directory = self._get_memory_directory(app_name, user_id)
+        return directory / f"{memory_id}.yaml"
+
+    def _serialize_content(self, content: types.Content) -> dict:
+        """Serialize Content object to dictionary."""
+        try:
+            return content.to_json_dict()
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"Failed to serialize content: {e}")
+
+    def _deserialize_content(self, content_dict: dict) -> types.Content:
+        """Deserialize Content object from dictionary."""
+        try:
+            return types.Content(**content_dict)
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"Failed to deserialize content: {e}")
+
+    def _extract_text_from_content(self, content: types.Content) -> str:
+        """Extract text content from a Content object for storage and search."""
+        if not content or not content.parts:
+            return ""
+        
+        text_parts = []
+        for part in content.parts:
+            if part.text:
+                text_parts.append(part.text)
+        
+        return " ".join(text_parts)
+
+    def _extract_search_terms(self, text: str) -> List[str]:
+        """Extract search terms from text content."""
+        # Extract words from text and convert to lowercase
+        words = re.findall(r'[A-Za-z]+', text.lower())
+        # Return unique words as a list
+        return sorted(set(words))
+
+    async def _add_session_to_memory_impl(self, session: "Session") -> None:
+        """Implementation of adding a session to memory."""
+        try:
+            # Add each event in the session as a separate memory entry
+            for event in session.events:
+                if not event.content or not event.content.parts:
+                    continue
+                
+                # Extract text content and search terms
+                text_content = self._extract_text_from_content(event.content)
+                search_terms = self._extract_search_terms(text_content)
+                
+                # Generate memory ID
+                memory_id = f"{session.id}_{event.timestamp}"
+                
+                # Create memory entry
+                memory_entry = {
+                    "id": memory_id,
+                    "app_name": session.app_name,
+                    "user_id": session.user_id,
+                    "content": self._serialize_content(event.content),
+                    "author": event.author,
+                    "timestamp": event.timestamp,
+                    "text_content": text_content,
+                    "search_terms": search_terms
+                }
+                
+                # Save to YAML file
+                file_path = self._get_memory_file_path(session.app_name, session.user_id, memory_id)
+                with open(file_path, 'w') as f:
+                    yaml.dump(memory_entry, f, default_flow_style=False, allow_unicode=True)
+                
+        except Exception as e:
+            raise RuntimeError(f"Failed to add session to memory: {e}")
+
+    async def _search_memory_impl(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> "SearchMemoryResponse":
+        """Implementation of searching memory."""
+        from google.adk.memory.base_memory_service import SearchMemoryResponse
+        from google.adk.memory.memory_entry import MemoryEntry
+        
+        try:
+            # Extract search terms from query
+            query_terms = self._extract_search_terms(query)
+            
+            if not query_terms:
+                # If no searchable terms in query, return empty response
+                return SearchMemoryResponse(memories=[])
+            
+            # Get memory directory for this user
+            memory_directory = self._get_memory_directory(app_name, user_id)
+            
+            # Find all memory files for this user
+            memory_files = list(memory_directory.glob("*.yaml"))
+            
+            # Filter entries that match any of the query terms
+            matching_memories = []
+            for file_path in memory_files:
+                try:
+                    with open(file_path, 'r') as f:
+                        entry = yaml.safe_load(f)
+                    
+                    # Check if any query term matches the search terms in this entry
+                    entry_search_terms = entry.get("search_terms", [])
+                    if any(term in entry_search_terms for term in query_terms):
+                        matching_memories.append(entry)
+                except (yaml.YAMLError, IOError):
+                    # Skip invalid files
+                    continue
+            
+            # Convert to MemoryEntry objects
+            memories = []
+            for entry in matching_memories:
+                content = self._deserialize_content(entry["content"])
+                # Format timestamp as ISO string
+                timestamp_str = None
+                if entry.get("timestamp"):
+                    timestamp_str = datetime.fromtimestamp(entry["timestamp"]).isoformat()
+                
+                memory_entry = MemoryEntry(
+                    content=content,
+                    author=entry.get("author"),
+                    timestamp=timestamp_str
+                )
+                memories.append(memory_entry)
+            
+            return SearchMemoryResponse(memories=memories)
+        except Exception as e:
+            raise RuntimeError(f"Failed to search memory: {e}")

google_adk_extras/sessions/__init__.py

@@ -0,0 +1,13 @@
+"""Custom session service implementations for Google ADK."""
+
+from .sql_session_service import SQLSessionService
+from .mongo_session_service import MongoSessionService
+from .redis_session_service import RedisSessionService
+from .yaml_file_session_service import YamlFileSessionService
+
+__all__ = [
+    "SQLSessionService",
+    "MongoSessionService",
+    "RedisSessionService",
+    "YamlFileSessionService",
+]

google_adk_extras/sessions/base_custom_session_service.py

@@ -0,0 +1,183 @@
+"""Base class for custom session services."""
+
+import abc
+from typing import Any, Optional
+
+from google.adk.sessions.base_session_service import BaseSessionService
+from google.adk.sessions.session import Session
+from google.adk.events.event import Event
+from google.adk.sessions.base_session_service import GetSessionConfig, ListSessionsResponse
+
+
+class BaseCustomSessionService(BaseSessionService, abc.ABC):
+    """Base class for custom session services with common functionality."""
+
+    def __init__(self):
+        """Initialize the base custom session service."""
+        super().__init__()
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """Initialize the session service.
+        
+        This method should be called before using the service to ensure
+        any required setup (database connections, etc.) is complete.
+        """
+        if not self._initialized:
+            await self._initialize_impl()
+            self._initialized = True
+
+    @abc.abstractmethod
+    async def _initialize_impl(self) -> None:
+        """Implementation of service initialization.
+        
+        This method should handle any setup required for the service to function,
+        such as database connections, creating tables, etc.
+        """
+        pass
+
+    async def cleanup(self) -> None:
+        """Clean up resources used by the session service.
+        
+        This method should be called when the service is no longer needed
+        to ensure proper cleanup of resources.
+        """
+        if self._initialized:
+            await self._cleanup_impl()
+            self._initialized = False
+
+    @abc.abstractmethod
+    async def _cleanup_impl(self) -> None:
+        """Implementation of service cleanup.
+        
+        This method should handle any cleanup required for the service,
+        such as closing database connections.
+        """
+        pass
+
+    async def create_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        """Create a new session."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._create_session_impl(
+            app_name=app_name,
+            user_id=user_id,
+            state=state,
+            session_id=session_id,
+        )
+
+    async def get_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        config: Optional[GetSessionConfig] = None,
+    ) -> Optional[Session]:
+        """Get a session by ID."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._get_session_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            config=config,
+        )
+
+    async def list_sessions(
+        self,
+        *,
+        app_name: str,
+        user_id: str
+    ) -> ListSessionsResponse:
+        """List all sessions for a user."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._list_sessions_impl(
+            app_name=app_name,
+            user_id=user_id,
+        )
+
+    async def delete_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str
+    ) -> None:
+        """Delete a session."""
+        if not self._initialized:
+            await self.initialize()
+        await self._delete_session_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+        )
+
+    async def append_event(self, session: Session, event: Event) -> Event:
+        """Append an event to a session."""
+        if not self._initialized:
+            await self.initialize()
+        # Update the session object
+        await super().append_event(session=session, event=event)
+        session.last_update_time = event.timestamp
+        # Update the storage
+        await self._append_event_impl(session=session, event=event)
+        return event
+
+    @abc.abstractmethod
+    async def _create_session_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        """Implementation of session creation."""
+        pass
+
+    @abc.abstractmethod
+    async def _get_session_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        config: Optional[GetSessionConfig] = None,
+    ) -> Optional[Session]:
+        """Implementation of session retrieval."""
+        pass
+
+    @abc.abstractmethod
+    async def _list_sessions_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str
+    ) -> ListSessionsResponse:
+        """Implementation of session listing."""
+        pass
+
+    @abc.abstractmethod
+    async def _delete_session_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str
+    ) -> None:
+        """Implementation of session deletion."""
+        pass
+
+    @abc.abstractmethod
+    async def _append_event_impl(self, session: Session, event: Event) -> None:
+        """Implementation of event appending."""
+        pass

google_adk_extras/sessions/mongo_session_service.py

@@ -0,0 +1,243 @@
+"""MongoDB-based session service implementation."""
+
+import json
+import time
+import uuid
+from typing import Any, Optional, Dict
+
+try:
+    from pymongo import MongoClient
+    from pymongo.errors import PyMongoError
+except ImportError:
+    raise ImportError(
+        "PyMongo is required for MongoSessionService. "
+        "Install it with: pip install pymongo"
+    )
+
+from google.adk.sessions.session import Session
+from google.adk.events.event import Event
+from google.adk.sessions.base_session_service import GetSessionConfig, ListSessionsResponse
+
+from .base_custom_session_service import BaseCustomSessionService
+
+
+class MongoSessionService(BaseCustomSessionService):
+    """MongoDB-based session service implementation."""
+
+    def __init__(self, connection_string: str, database_name: str = "adk_sessions"):
+        """Initialize the MongoDB session service.
+        
+        Args:
+            connection_string: MongoDB connection string
+            database_name: Name of the database to use
+        """
+        super().__init__()
+        self.connection_string = connection_string
+        self.database_name = database_name
+        self.client: Optional[MongoClient] = None
+        self.db = None
+        self.collection = None
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the MongoDB connection."""
+        try:
+            self.client = MongoClient(self.connection_string)
+            self.db = self.client[self.database_name]
+            self.collection = self.db.sessions
+            
+            # Create indexes for better performance
+            self.collection.create_index([("app_name", 1), ("user_id", 1)])
+            self.collection.create_index("id")
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to initialize MongoDB session service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up MongoDB connections."""
+        if self.client:
+            self.client.close()
+            self.client = None
+        self.db = None
+        self.collection = None
+
+    def _serialize_events(self, events: list[Event]) -> list[Dict]:
+        """Serialize events to dictionaries."""
+        return [event.model_dump() for event in events]
+
+    def _deserialize_events(self, event_dicts: list[Dict]) -> list[Event]:
+        """Deserialize events from dictionaries."""
+        return [Event(**event_dict) for event_dict in event_dicts]
+
+    async def _create_session_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        """Implementation of session creation."""
+        try:
+            # Generate session ID if not provided
+            session_id = session_id or str(uuid.uuid4())
+            
+            # Create session object
+            session = Session(
+                id=session_id,
+                app_name=app_name,
+                user_id=user_id,
+                state=state or {},
+                events=[],
+                last_update_time=time.time()
+            )
+            
+            # Create document for MongoDB
+            document = {
+                "_id": session_id,
+                "id": session_id,
+                "app_name": app_name,
+                "user_id": user_id,
+                "state": session.state,
+                "events": self._serialize_events(session.events),
+                "last_update_time": session.last_update_time
+            }
+            
+            # Insert into MongoDB
+            self.collection.insert_one(document)
+            
+            return session
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to create session: {e}")
+
+    async def _get_session_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        config: Optional[GetSessionConfig] = None,
+    ) -> Optional[Session]:
+        """Implementation of session retrieval."""
+        try:
+            # Retrieve from MongoDB
+            document = self.collection.find_one({
+                "_id": session_id,
+                "app_name": app_name,
+                "user_id": user_id
+            })
+            
+            if not document:
+                return None
+            
+            # Create session object
+            session = Session(
+                id=document["id"],
+                app_name=document["app_name"],
+                user_id=document["user_id"],
+                state=document["state"],
+                events=self._deserialize_events(document["events"]),
+                last_update_time=document["last_update_time"]
+            )
+            
+            # Apply config filters if provided
+            if config:
+                if config.num_recent_events:
+                    session.events = session.events[-config.num_recent_events:]
+                if config.after_timestamp:
+                    filtered_events = [
+                        event for event in session.events 
+                        if event.timestamp >= config.after_timestamp
+                    ]
+                    session.events = filtered_events
+            
+            return session
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to get session: {e}")
+
+    async def _list_sessions_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str
+    ) -> ListSessionsResponse:
+        """Implementation of session listing."""
+        try:
+            # Retrieve all sessions for user (without events for performance)
+            cursor = self.collection.find(
+                {
+                    "app_name": app_name,
+                    "user_id": user_id
+                },
+                {
+                    "_id": 1,
+                    "id": 1,
+                    "app_name": 1,
+                    "user_id": 1,
+                    "state": 1,
+                    "last_update_time": 1
+                }
+            )
+            
+            # Create session objects without events
+            sessions = []
+            for document in cursor:
+                session = Session(
+                    id=document["id"],
+                    app_name=document["app_name"],
+                    user_id=document["user_id"],
+                    state=document["state"],
+                    events=[],  # Empty events for listing
+                    last_update_time=document["last_update_time"]
+                )
+                sessions.append(session)
+            
+            return ListSessionsResponse(sessions=sessions)
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to list sessions: {e}")
+
+    async def _delete_session_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str
+    ) -> None:
+        """Implementation of session deletion."""
+        try:
+            # Delete from MongoDB
+            self.collection.delete_one({
+                "_id": session_id,
+                "app_name": app_name,
+                "user_id": user_id
+            })
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to delete session: {e}")
+
+    async def _append_event_impl(self, session: Session, event: Event) -> None:
+        """Implementation of event appending."""
+        try:
+            # Prepare update data
+            update_data = {
+                "$set": {
+                    "events": self._serialize_events(session.events),
+                    "last_update_time": session.last_update_time
+                }
+            }
+            
+            # Apply state changes from event if present
+            if event.actions and event.actions.state_delta:
+                update_data["$set"]["state"] = session.state
+            
+            # Update session in MongoDB
+            result = self.collection.update_one(
+                {
+                    "_id": session.id,
+                    "app_name": session.app_name,
+                    "user_id": session.user_id
+                },
+                update_data
+            )
+            
+            if result.matched_count == 0:
+                raise ValueError(f"Session {session.id} not found")
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to append event: {e}")