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

google_adk_extras/memory/mongo_memory_service.py

@@ -0,0 +1,174 @@
+"""MongoDB-based memory service implementation using PyMongo."""
+
+import json
+import logging
+from typing import Optional, List
+import re
+from datetime import datetime
+
+try:
+    from pymongo import MongoClient
+    from pymongo.errors import PyMongoError
+    from bson import ObjectId
+except ImportError:
+    raise ImportError(
+        "PyMongo is required for MongoMemoryService. "
+        "Install it with: pip install pymongo"
+    )
+
+from google.genai import types
+from .base_custom_memory_service import BaseCustomMemoryService
+
+
+logger = logging.getLogger('google_adk_extras.' + __name__)
+
+
+class MongoMemoryService(BaseCustomMemoryService):
+    """MongoDB-based memory service implementation."""
+
+    def __init__(self, connection_string: str, database_name: str = "adk_memory"):
+        """Initialize the MongoDB memory 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: Optional[object] = 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]
+            
+            # Create indexes for efficient querying
+            memory_collection = self.db["memory_entries"]
+            memory_collection.create_index([("app_name", 1), ("user_id", 1)])
+            memory_collection.create_index([("search_terms", 1)])
+            memory_collection.create_index([("app_name", 1), ("user_id", 1), ("search_terms", 1)])
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to initialize MongoDB memory service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up MongoDB connections."""
+        if self.client:
+            self.client.close()
+            self.client = None
+        self.db = None
+
+    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."""
+        if not self.db:
+            raise RuntimeError("Service not initialized")
+        
+        try:
+            memory_collection = self.db["memory_entries"]
+            
+            # 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)
+                
+                # Create memory document
+                memory_document = {
+                    "app_name": session.app_name,
+                    "user_id": session.user_id,
+                    "content": self._serialize_content(event.content),
+                    "author": event.author,
+                    "timestamp": datetime.fromtimestamp(event.timestamp) if event.timestamp else None,
+                    "text_content": text_content,
+                    "search_terms": search_terms
+                }
+                
+                # Save to MongoDB
+                memory_collection.insert_one(memory_document)
+                
+        except PyMongoError 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
+        
+        if not self.db:
+            raise RuntimeError("Service not initialized")
+        
+        try:
+            memory_collection = self.db["memory_entries"]
+            
+            # 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=[])
+            
+            # Build query - find entries that match any of the query words
+            # Using MongoDB's $in operator to match any of the search terms
+            mongo_query = {
+                "app_name": app_name,
+                "user_id": user_id,
+                "search_terms": {"$in": query_terms}
+            }
+            
+            # Execute query
+            cursor = memory_collection.find(mongo_query)
+            
+            # Convert to MemoryEntry objects
+            memories = []
+            for doc in cursor:
+                content = self._deserialize_content(doc["content"])
+                memory_entry = MemoryEntry(
+                    content=content,
+                    author=doc.get("author"),
+                    timestamp=doc["timestamp"].isoformat() if doc.get("timestamp") else None
+                )
+                memories.append(memory_entry)
+            
+            return SearchMemoryResponse(memories=memories)
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to search memory: {e}")

google_adk_extras/memory/redis_memory_service.py

@@ -0,0 +1,188 @@
+"""Redis-based memory service implementation using redis-py."""
+
+import json
+import logging
+from typing import Optional, List
+import re
+from datetime import datetime
+
+try:
+    import redis
+    from redis.exceptions import RedisError
+except ImportError:
+    raise ImportError(
+        "Redis-py is required for RedisMemoryService. "
+        "Install it with: pip install redis"
+    )
+
+from google.genai import types
+from .base_custom_memory_service import BaseCustomMemoryService
+
+
+logger = logging.getLogger('google_adk_extras.' + __name__)
+
+
+class RedisMemoryService(BaseCustomMemoryService):
+    """Redis-based memory service implementation."""
+
+    def __init__(self, host: str = "localhost", port: int = 6379, db: int = 0):
+        """Initialize the Redis memory service.
+        
+        Args:
+            host: Redis server host
+            port: Redis server port
+            db: Redis database number
+        """
+        super().__init__()
+        self.host = host
+        self.port = port
+        self.db = db
+        self.client: Optional[redis.Redis] = None
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the Redis connection."""
+        try:
+            self.client = redis.Redis(host=self.host, port=self.port, db=self.db)
+            # Test connection
+            self.client.ping()
+        except RedisError as e:
+            raise RuntimeError(f"Failed to initialize Redis memory service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up Redis connections."""
+        if self.client:
+            self.client.close()
+            self.client = None
+
+    def _serialize_content(self, content: types.Content) -> str:
+        """Serialize Content object to JSON string."""
+        try:
+            return json.dumps(content.to_json_dict())
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"Failed to serialize content: {e}")
+
+    def _deserialize_content(self, content_str: str) -> types.Content:
+        """Deserialize Content object from JSON string."""
+        try:
+            content_dict = json.loads(content_str) if content_str else {}
+            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))
+
+    def _get_user_key(self, app_name: str, user_id: str) -> str:
+        """Generate Redis key for user's memory entries."""
+        return f"memory:{app_name}:{user_id}"
+
+    async def _add_session_to_memory_impl(self, session: "Session") -> None:
+        """Implementation of adding a session to memory."""
+        if not self.client:
+            raise RuntimeError("Service not initialized")
+        
+        try:
+            user_key = self._get_user_key(session.app_name, session.user_id)
+            
+            # 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)
+                
+                # Create memory entry
+                memory_entry = {
+                    "id": f"{session.id}:{event.timestamp}",
+                    "content": self._serialize_content(event.content),
+                    "author": event.author,
+                    "timestamp": event.timestamp,
+                    "text_content": text_content,
+                    "search_terms": search_terms
+                }
+                
+                # Store in Redis with a score based on timestamp for ordering
+                score = event.timestamp if event.timestamp else 0
+                self.client.zadd(user_key, {json.dumps(memory_entry): score})
+                
+        except RedisError as e:
+            raise RuntimeError(f"Failed to add session to memory: {e}")
+        except (TypeError, ValueError) as e:
+            raise RuntimeError(f"Failed to serialize memory entry: {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
+        
+        if not self.client:
+            raise RuntimeError("Service not initialized")
+        
+        try:
+            user_key = self._get_user_key(app_name, user_id)
+            
+            # 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 all memory entries for the user
+            memory_entries = self.client.zrange(user_key, 0, -1, withscores=True)
+            
+            # Filter entries that match any of the query terms
+            matching_memories = []
+            for entry_data, _ in memory_entries:
+                try:
+                    entry = json.loads(entry_data)
+                    
+                    # 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 (TypeError, ValueError):
+                    # Skip invalid entries
+                    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 RedisError as e:
+            raise RuntimeError(f"Failed to search memory: {e}")
+        except (TypeError, ValueError) as e:
+            raise RuntimeError(f"Failed to deserialize memory entry: {e}")

google_adk_extras/memory/sql_memory_service.py

@@ -0,0 +1,213 @@
+"""SQL-based memory service implementation using SQLAlchemy."""
+
+import json
+import logging
+from typing import Optional, List
+import re
+from datetime import datetime, timezone
+
+try:
+    from sqlalchemy import create_engine, Column, String, Text, DateTime, Integer, Index
+    from sqlalchemy.orm import declarative_base, sessionmaker
+    from sqlalchemy.exc import SQLAlchemyError
+except ImportError:
+    raise ImportError(
+        "SQLAlchemy is required for SQLMemoryService. "
+        "Install it with: pip install sqlalchemy"
+    )
+
+from google.genai import types
+from .base_custom_memory_service import BaseCustomMemoryService
+
+
+logger = logging.getLogger('google_adk_extras.' + __name__)
+
+# Use the modern declarative_base import
+Base = declarative_base()
+
+
+class SQLMemoryModel(Base):
+    """SQLAlchemy model for storing memory entries."""
+    __tablename__ = 'adk_memory_entries'
+
+    # Auto-incrementing ID
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    
+    # Identifiers
+    app_name = Column(String, nullable=False, index=True)
+    user_id = Column(String, nullable=False, index=True)
+    
+    # Content data
+    content_json = Column(Text, nullable=False)  # JSON string of Content object
+    author = Column(String, nullable=True)
+    timestamp = Column(DateTime(timezone=True), nullable=True)
+    
+    # Search optimization
+    text_content = Column(Text, nullable=False)  # Extracted text for searching
+    search_terms = Column(Text, nullable=False)  # Space-separated search terms
+
+    # Composite index for efficient querying
+    __table_args__ = (
+        Index('idx_app_user_query', 'app_name', 'user_id', 'search_terms'),
+    )
+
+
+class SQLMemoryService(BaseCustomMemoryService):
+    """SQL-based memory service implementation."""
+
+    def __init__(self, database_url: str):
+        """Initialize the SQL memory service.
+        
+        Args:
+            database_url: Database connection string (e.g., 'sqlite:///memory.db')
+        """
+        super().__init__()
+        self.database_url = database_url
+        self.engine: Optional[object] = None
+        self.session_local: Optional[object] = None
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the database connection and create tables."""
+        try:
+            self.engine = create_engine(self.database_url)
+            Base.metadata.create_all(self.engine)
+            self.session_local = sessionmaker(
+                autocommit=False, 
+                autoflush=False, 
+                bind=self.engine
+            )
+        except SQLAlchemyError as e:
+            raise RuntimeError(f"Failed to initialize SQL memory service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up database connections."""
+        if self.engine:
+            self.engine.dispose()
+            self.engine = None
+        self.session_local = None
+
+    def _get_db_session(self):
+        """Get a database session."""
+        if not self.session_local:
+            raise RuntimeError("Service not initialized")
+        return self.session_local()
+
+    def _serialize_content(self, content: types.Content) -> str:
+        """Serialize Content object to JSON string."""
+        try:
+            return json.dumps(content.to_json_dict())
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"Failed to serialize content: {e}")
+
+    def _deserialize_content(self, content_str: str) -> types.Content:
+        """Deserialize Content object from JSON string."""
+        try:
+            content_dict = json.loads(content_str) if content_str else {}
+            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) -> 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 space-separated unique words
+        return " ".join(sorted(set(words)))
+
+    async def _add_session_to_memory_impl(self, session: "Session") -> None:
+        """Implementation of adding a session to memory."""
+        db_session = self._get_db_session()
+        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)
+                
+                # Create memory model
+                db_memory = SQLMemoryModel(
+                    app_name=session.app_name,
+                    user_id=session.user_id,
+                    content_json=self._serialize_content(event.content),
+                    author=event.author,
+                    timestamp=datetime.fromtimestamp(event.timestamp, tz=timezone.utc) if event.timestamp else None,
+                    text_content=text_content,
+                    search_terms=search_terms
+                )
+                
+                # Save to database
+                db_session.add(db_memory)
+            
+            db_session.commit()
+        except SQLAlchemyError as e:
+            db_session.rollback()
+            raise RuntimeError(f"Failed to add session to memory: {e}")
+        finally:
+            db_session.close()
+
+    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
+        
+        db_session = self._get_db_session()
+        try:
+            # Extract search terms from query
+            query_terms = self._extract_search_terms(query)
+            query_words = query_terms.split() if query_terms else []
+            
+            if not query_words:
+                # If no searchable terms in query, return empty response
+                return SearchMemoryResponse(memories=[])
+            
+            # Build query - find entries that match any of the query words
+            db_query = db_session.query(SQLMemoryModel).filter(
+                SQLMemoryModel.app_name == app_name,
+                SQLMemoryModel.user_id == user_id
+            )
+            
+            # Add search term filters - use OR logic to match any query word
+            from sqlalchemy import or_
+            conditions = []
+            for word in query_words:
+                conditions.append(SQLMemoryModel.search_terms.contains(word))
+            
+            if conditions:
+                db_query = db_query.filter(or_(*conditions))
+            
+            # Execute query
+            db_memories = db_query.all()
+            
+            # Convert to MemoryEntry objects
+            memories = []
+            for db_memory in db_memories:
+                content = self._deserialize_content(db_memory.content_json)
+                memory_entry = MemoryEntry(
+                    content=content,
+                    author=db_memory.author,
+                    timestamp=db_memory.timestamp.isoformat() if db_memory.timestamp else None
+                )
+                memories.append(memory_entry)
+            
+            return SearchMemoryResponse(memories=memories)
+        except SQLAlchemyError as e:
+            raise RuntimeError(f"Failed to search memory: {e}")
+        finally:
+            db_session.close()