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

google_adk_extras/memory/sql_memory_service.py

@@ -53,7 +53,12 @@ class SQLMemoryModel(Base):
 
 
 class SQLMemoryService(BaseCustomMemoryService):
-    """SQL-based memory service implementation."""
+    """SQL-based memory service implementation.
+
+    This service stores memory entries in a SQL database using SQLAlchemy.
+    It supports efficient searching of memory entries by extracting and indexing
+    text content from conversation events.
+    """
 
     def __init__(self, database_url: str):
         """Initialize the SQL memory service.
@@ -67,7 +72,11 @@ class SQLMemoryService(BaseCustomMemoryService):
         self.session_local: Optional[object] = None
 
     async def _initialize_impl(self) -> None:
-        """Initialize the database connection and create tables."""
+        """Initialize the database connection and create tables.
+        
+        Raises:
+            RuntimeError: If database initialization fails.
+        """
         try:
             self.engine = create_engine(self.database_url)
             Base.metadata.create_all(self.engine)
@@ -87,20 +96,47 @@ class SQLMemoryService(BaseCustomMemoryService):
         self.session_local = None
 
     def _get_db_session(self):
-        """Get a database session."""
+        """Get a database session.
+        
+        Returns:
+            A database session object.
+            
+        Raises:
+            RuntimeError: If the service is not initialized.
+        """
         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."""
+        """Serialize Content object to JSON string.
+        
+        Args:
+            content: The Content object to serialize.
+            
+        Returns:
+            JSON string representation of the content.
+            
+        Raises:
+            ValueError: If serialization fails.
+        """
         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."""
+        """Deserialize Content object from JSON string.
+        
+        Args:
+            content_str: JSON string representation of the content.
+            
+        Returns:
+            The deserialized Content object.
+            
+        Raises:
+            ValueError: If deserialization fails.
+        """
         try:
             content_dict = json.loads(content_str) if content_str else {}
             return types.Content(**content_dict)
@@ -108,7 +144,14 @@ class SQLMemoryService(BaseCustomMemoryService):
             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."""
+        """Extract text content from a Content object for storage and search.
+        
+        Args:
+            content: The Content object to extract text from.
+            
+        Returns:
+            Extracted text content.
+        """
         if not content or not content.parts:
             return ""
         
@@ -120,14 +163,28 @@ class SQLMemoryService(BaseCustomMemoryService):
         return " ".join(text_parts)
 
     def _extract_search_terms(self, text: str) -> str:
-        """Extract search terms from text content."""
+        """Extract search terms from text content.
+        
+        Args:
+            text: The text to extract search terms from.
+            
+        Returns:
+            Space-separated unique search terms.
+        """
         # 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."""
+        """Implementation of adding a session to memory.
+        
+        Args:
+            session: The session to add to memory.
+            
+        Raises:
+            RuntimeError: If adding the session to memory fails.
+        """
         db_session = self._get_db_session()
         try:
             # Add each event in the session as a separate memory entry
@@ -163,9 +220,25 @@ class SQLMemoryService(BaseCustomMemoryService):
     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
+        """Implementation of searching memory.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The id of the user.
+            query: The query to search for.
+            
+        Returns:
+            A SearchMemoryResponse containing the matching memories.
+            
+        Raises:
+            RuntimeError: If searching memory fails.
+        """
+        try:
+            from google.adk.memory.base_memory_service import SearchMemoryResponse
+            from google.adk.memory.memory_entry import MemoryEntry
+        except Exception:  # Fallback when ADK surface differs
+            from types import SimpleNamespace as MemoryEntry  # type: ignore
+            SearchMemoryResponse = None  # type: ignore
         
         db_session = self._get_db_session()
         try:
@@ -199,15 +272,28 @@ class SQLMemoryService(BaseCustomMemoryService):
             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
-                )
+                try:
+                    memory_entry = MemoryEntry(
+                        content=content,
+                        author=db_memory.author,
+                        timestamp=db_memory.timestamp.isoformat() if db_memory.timestamp else None
+                    )
+                except TypeError:
+                    # SimpleNamespace fallback
+                    from types import SimpleNamespace
+                    memory_entry = SimpleNamespace(
+                        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)
+
+            if SearchMemoryResponse is not None:
+                return SearchMemoryResponse(memories=memories)
+            else:
+                from types import SimpleNamespace
+                return SimpleNamespace(memories=memories)
         except SQLAlchemyError as e:
             raise RuntimeError(f"Failed to search memory: {e}")
         finally:
-            db_session.close()
+            db_session.close()

google_adk_extras/memory/yaml_file_memory_service.py

@@ -18,13 +18,18 @@ logger = logging.getLogger('google_adk_extras.' + __name__)
 
 
 class YamlFileMemoryService(BaseCustomMemoryService):
-    """YAML file-based memory service implementation."""
+    """YAML file-based memory service implementation.
+
+    This service stores memory entries in YAML files in a hierarchical directory structure.
+    Each memory entry is stored in a separate YAML file organized by app name and user ID.
+    Memory entries are searchable by extracting and indexing text content from conversation events.
+    """
 
     def __init__(self, base_directory: str = "./memory"):
         """Initialize the YAML file memory service.
         
         Args:
-            base_directory: Base directory for storing memory files
+            base_directory: Base directory for storing memory files. Defaults to "./memory".
         """
         super().__init__()
         self.base_directory = Path(base_directory)
@@ -32,7 +37,10 @@ class YamlFileMemoryService(BaseCustomMemoryService):
         self.base_directory.mkdir(parents=True, exist_ok=True)
 
     async def _initialize_impl(self) -> None:
-        """Initialize the file system memory service."""
+        """Initialize the file system memory service.
+        
+        Ensures the base directory exists.
+        """
         # Ensure base directory exists
         self.base_directory.mkdir(parents=True, exist_ok=True)
 
@@ -41,32 +49,76 @@ class YamlFileMemoryService(BaseCustomMemoryService):
         pass
 
     def _get_memory_directory(self, app_name: str, user_id: str) -> Path:
-        """Generate directory path for memory entries."""
+        """Generate directory path for memory entries.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            
+        Returns:
+            Path to the memory directory.
+        """
         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."""
+        """Generate file path for a memory entry.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            memory_id: The ID of the memory entry.
+            
+        Returns:
+            Path to the memory file.
+        """
         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."""
+        """Serialize Content object to dictionary.
+        
+        Args:
+            content: The Content object to serialize.
+            
+        Returns:
+            Dictionary representation of the content.
+            
+        Raises:
+            ValueError: If serialization fails.
+        """
         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."""
+        """Deserialize Content object from dictionary.
+        
+        Args:
+            content_dict: Dictionary representation of the content.
+            
+        Returns:
+            The deserialized Content object.
+            
+        Raises:
+            ValueError: If deserialization fails.
+        """
         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."""
+        """Extract text content from a Content object for storage and search.
+        
+        Args:
+            content: The Content object to extract text from.
+            
+        Returns:
+            Extracted text content.
+        """
         if not content or not content.parts:
             return ""
         
@@ -78,14 +130,28 @@ class YamlFileMemoryService(BaseCustomMemoryService):
         return " ".join(text_parts)
 
     def _extract_search_terms(self, text: str) -> List[str]:
-        """Extract search terms from text content."""
+        """Extract search terms from text content.
+        
+        Args:
+            text: The text to extract search terms from.
+            
+        Returns:
+            List of unique search terms.
+        """
         # 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."""
+        """Implementation of adding a session to memory.
+        
+        Args:
+            session: The session to add to memory.
+            
+        Raises:
+            RuntimeError: If adding the session to memory fails.
+        """
         try:
             # Add each event in the session as a separate memory entry
             for event in session.events:
@@ -122,9 +188,25 @@ class YamlFileMemoryService(BaseCustomMemoryService):
     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
+        """Implementation of searching memory.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The id of the user.
+            query: The query to search for.
+            
+        Returns:
+            A SearchMemoryResponse containing the matching memories.
+            
+        Raises:
+            RuntimeError: If searching memory fails.
+        """
+        try:
+            from google.adk.memory.base_memory_service import SearchMemoryResponse
+            from google.adk.memory.memory_entry import MemoryEntry
+        except Exception:
+            from types import SimpleNamespace as MemoryEntry  # type: ignore
+            SearchMemoryResponse = None  # type: ignore
         
         try:
             # Extract search terms from query
@@ -163,14 +245,25 @@ class YamlFileMemoryService(BaseCustomMemoryService):
                 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
-                )
+                try:
+                    memory_entry = MemoryEntry(
+                        content=content,
+                        author=entry.get("author"),
+                        timestamp=timestamp_str
+                    )
+                except TypeError:
+                    from types import SimpleNamespace
+                    memory_entry = SimpleNamespace(
+                        content=content,
+                        author=entry.get("author"),
+                        timestamp=timestamp_str
+                    )
                 memories.append(memory_entry)
-            
-            return SearchMemoryResponse(memories=memories)
+
+            if SearchMemoryResponse is not None:
+                return SearchMemoryResponse(memories=memories)
+            else:
+                from types import SimpleNamespace
+                return SimpleNamespace(memories=memories)
         except Exception as e:
-            raise RuntimeError(f"Failed to search memory: {e}")
+            raise RuntimeError(f"Failed to search memory: {e}")

google_adk_extras/sessions/__init__.py

@@ -1,13 +1,29 @@
-"""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",
-]
+"""Custom session service implementations for Google ADK.
+
+Optional backends are imported lazily based on installed dependencies.
+"""
+
+try:
+    from .sql_session_service import SQLSessionService  # type: ignore
+except Exception:
+    SQLSessionService = None  # type: ignore
+
+try:
+    from .mongo_session_service import MongoSessionService  # type: ignore
+except Exception:
+    MongoSessionService = None  # type: ignore
+
+try:
+    from .redis_session_service import RedisSessionService  # type: ignore
+except Exception:
+    RedisSessionService = None  # type: ignore
+
+try:
+    from .yaml_file_session_service import YamlFileSessionService  # type: ignore
+except Exception:
+    YamlFileSessionService = None  # type: ignore
+
+__all__ = []
+for _name in ("SQLSessionService", "MongoSessionService", "RedisSessionService", "YamlFileSessionService"):
+    if globals().get(_name) is not None:
+        __all__.append(_name)

google_adk_extras/sessions/base_custom_session_service.py

@@ -10,7 +10,11 @@ from google.adk.sessions.base_session_service import GetSessionConfig, ListSessi
 
 
 class BaseCustomSessionService(BaseSessionService, abc.ABC):
-    """Base class for custom session services with common functionality."""
+    """Base class for custom session services with common functionality.
+
+    This abstract base class provides a foundation for implementing custom
+    session services with automatic initialization and cleanup handling.
+    """
 
     def __init__(self):
         """Initialize the base custom session service."""
@@ -22,6 +26,9 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         
         This method should be called before using the service to ensure
         any required setup (database connections, etc.) is complete.
+        
+        Raises:
+            RuntimeError: If initialization fails.
         """
         if not self._initialized:
             await self._initialize_impl()
@@ -33,6 +40,9 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         
         This method should handle any setup required for the service to function,
         such as database connections, creating tables, etc.
+        
+        Raises:
+            RuntimeError: If initialization fails.
         """
         pass
 
@@ -63,7 +73,21 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         state: Optional[dict[str, Any]] = None,
         session_id: Optional[str] = None,
     ) -> Session:
-        """Create a new session."""
+        """Create a new session.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            state: Optional initial state for the session.
+            session_id: Optional specific ID for the session. If not provided,
+                a UUID will be generated.
+                
+        Returns:
+            The created Session object.
+            
+        Raises:
+            RuntimeError: If session creation fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._create_session_impl(
@@ -81,7 +105,20 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         session_id: str,
         config: Optional[GetSessionConfig] = None,
     ) -> Optional[Session]:
-        """Get a session by ID."""
+        """Get a session by ID.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session to retrieve.
+            config: Optional configuration for session retrieval.
+            
+        Returns:
+            The Session object if found, None otherwise.
+            
+        Raises:
+            RuntimeError: If session retrieval fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._get_session_impl(
@@ -97,7 +134,18 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         app_name: str,
         user_id: str
     ) -> ListSessionsResponse:
-        """List all sessions for a user."""
+        """List all sessions for a user.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            
+        Returns:
+            A ListSessionsResponse containing the sessions.
+            
+        Raises:
+            RuntimeError: If session listing fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._list_sessions_impl(
@@ -112,7 +160,16 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         user_id: str,
         session_id: str
     ) -> None:
-        """Delete a session."""
+        """Delete a session.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session to delete.
+            
+        Raises:
+            RuntimeError: If session deletion fails.
+        """
         if not self._initialized:
             await self.initialize()
         await self._delete_session_impl(
@@ -122,7 +179,18 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         )
 
     async def append_event(self, session: Session, event: Event) -> Event:
-        """Append an event to a session."""
+        """Append an event to a session.
+        
+        Args:
+            session: The session to append the event to.
+            event: The event to append.
+            
+        Returns:
+            The appended event.
+            
+        Raises:
+            RuntimeError: If appending the event fails.
+        """
         if not self._initialized:
             await self.initialize()
         # Update the session object
@@ -141,7 +209,20 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         state: Optional[dict[str, Any]] = None,
         session_id: Optional[str] = None,
     ) -> Session:
-        """Implementation of session creation."""
+        """Implementation of session creation.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            state: Optional initial state for the session.
+            session_id: Optional specific ID for the session.
+            
+        Returns:
+            The created Session object.
+            
+        Raises:
+            RuntimeError: If session creation fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -153,7 +234,20 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         session_id: str,
         config: Optional[GetSessionConfig] = None,
     ) -> Optional[Session]:
-        """Implementation of session retrieval."""
+        """Implementation of session retrieval.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session to retrieve.
+            config: Optional configuration for session retrieval.
+            
+        Returns:
+            The Session object if found, None otherwise.
+            
+        Raises:
+            RuntimeError: If session retrieval fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -163,7 +257,18 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         app_name: str,
         user_id: str
     ) -> ListSessionsResponse:
-        """Implementation of session listing."""
+        """Implementation of session listing.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            
+        Returns:
+            A ListSessionsResponse containing the sessions.
+            
+        Raises:
+            RuntimeError: If session listing fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -174,10 +279,27 @@ class BaseCustomSessionService(BaseSessionService, abc.ABC):
         user_id: str,
         session_id: str
     ) -> None:
-        """Implementation of session deletion."""
+        """Implementation of session deletion.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session to delete.
+            
+        Raises:
+            RuntimeError: If session deletion fails.
+        """
         pass
 
     @abc.abstractmethod
     async def _append_event_impl(self, session: Session, event: Event) -> None:
-        """Implementation of event appending."""
+        """Implementation of event appending.
+        
+        Args:
+            session: The session to append the event to.
+            event: The event to append.
+            
+        Raises:
+            RuntimeError: If appending the event fails.
+        """
         pass