memorisdk 2.0.1__py3-none-any.whl → 2.1.1__py3-none-any.whl

memori/database/adapters/mongodb_adapter.py

@@ -0,0 +1,739 @@
+"""
+MongoDB adapter for Memori memory storage
+Implements MongoDB-specific CRUD operations for memories
+"""
+
+import json
+from datetime import datetime, timezone
+from typing import Any
+from uuid import uuid4
+
+from loguru import logger
+
+try:
+    import pymongo  # noqa: F401
+    from bson import ObjectId  # noqa: F401
+    from pymongo.collection import Collection  # noqa: F401
+    from pymongo.errors import DuplicateKeyError, OperationFailure  # noqa: F401
+
+    PYMONGO_AVAILABLE = True
+except ImportError:
+    PYMONGO_AVAILABLE = False
+
+from ..connectors.mongodb_connector import MongoDBConnector
+
+
+class MongoDBAdapter:
+    """MongoDB-specific adapter for memory storage and retrieval"""
+
+    def __init__(self, connector: MongoDBConnector):
+        """Initialize MongoDB adapter"""
+        if not PYMONGO_AVAILABLE:
+            raise ImportError(
+                "pymongo is required for MongoDB support. Install with: pip install pymongo"
+            )
+
+        self.connector = connector
+        self.database = connector.get_database()
+
+        # Collection names
+        self.CHAT_HISTORY_COLLECTION = "chat_history"
+        self.SHORT_TERM_MEMORY_COLLECTION = "short_term_memory"
+        self.LONG_TERM_MEMORY_COLLECTION = "long_term_memory"
+
+        # Initialize collections
+        self._initialize_collections()
+
+    def _initialize_collections(self):
+        """Initialize MongoDB collections with proper indexes"""
+        try:
+            # Ensure collections exist
+            collections = [
+                self.CHAT_HISTORY_COLLECTION,
+                self.SHORT_TERM_MEMORY_COLLECTION,
+                self.LONG_TERM_MEMORY_COLLECTION,
+            ]
+
+            existing_collections = self.database.list_collection_names()
+            for collection_name in collections:
+                if collection_name not in existing_collections:
+                    self.database.create_collection(collection_name)
+                    logger.info(f"Created MongoDB collection: {collection_name}")
+
+            # Create basic indexes
+            self._create_indexes()
+
+        except Exception as e:
+            logger.warning(f"Failed to initialize MongoDB collections: {e}")
+
+    def _create_indexes(self):
+        """Create essential indexes for performance"""
+        try:
+            # Chat history indexes
+            chat_collection = self.database[self.CHAT_HISTORY_COLLECTION]
+            chat_collection.create_index([("chat_id", 1)], unique=True, background=True)
+            chat_collection.create_index(
+                [("namespace", 1), ("session_id", 1)], background=True
+            )
+            chat_collection.create_index([("timestamp", -1)], background=True)
+
+            # Short-term memory indexes
+            st_collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+            st_collection.create_index([("memory_id", 1)], unique=True, background=True)
+            st_collection.create_index(
+                [("namespace", 1), ("category_primary", 1), ("importance_score", -1)],
+                background=True,
+            )
+            st_collection.create_index([("expires_at", 1)], background=True)
+            st_collection.create_index([("created_at", -1)], background=True)
+            st_collection.create_index(
+                [("searchable_content", "text"), ("summary", "text")], background=True
+            )
+
+            # Long-term memory indexes
+            lt_collection = self.database[self.LONG_TERM_MEMORY_COLLECTION]
+            lt_collection.create_index([("memory_id", 1)], unique=True, background=True)
+            lt_collection.create_index(
+                [("namespace", 1), ("category_primary", 1), ("importance_score", -1)],
+                background=True,
+            )
+            lt_collection.create_index([("classification", 1)], background=True)
+            lt_collection.create_index([("topic", 1)], background=True)
+            lt_collection.create_index([("created_at", -1)], background=True)
+            lt_collection.create_index(
+                [("searchable_content", "text"), ("summary", "text")], background=True
+            )
+
+            logger.debug("MongoDB indexes created successfully")
+
+        except Exception as e:
+            logger.warning(f"Failed to create MongoDB indexes: {e}")
+
+    def _convert_memory_to_document(
+        self, memory_data: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Convert memory data to MongoDB document format"""
+        document = memory_data.copy()
+
+        # Ensure datetime fields are datetime objects
+        datetime_fields = [
+            "created_at",
+            "expires_at",
+            "last_accessed",
+            "extraction_timestamp",
+        ]
+        for field in datetime_fields:
+            if field in document and document[field] is not None:
+                if isinstance(document[field], str):
+                    try:
+                        document[field] = datetime.fromisoformat(
+                            document[field].replace("Z", "+00:00")
+                        )
+                    except:
+                        document[field] = datetime.now(timezone.utc)
+                elif not isinstance(document[field], datetime):
+                    document[field] = datetime.now(timezone.utc)
+
+        # Handle JSON fields that might be strings
+        json_fields = [
+            "processed_data",
+            "entities_json",
+            "keywords_json",
+            "supersedes_json",
+            "related_memories_json",
+            "metadata",
+        ]
+        for field in json_fields:
+            if field in document and isinstance(document[field], str):
+                try:
+                    document[field] = json.loads(document[field])
+                except:
+                    pass  # Keep as string if not valid JSON
+
+        # Ensure required fields have defaults
+        if "created_at" not in document:
+            document["created_at"] = datetime.now(timezone.utc)
+        if "importance_score" not in document:
+            document["importance_score"] = 0.5
+        if "access_count" not in document:
+            document["access_count"] = 0
+        if "namespace" not in document:
+            document["namespace"] = "default"
+
+        return document
+
+    def _convert_document_to_memory(self, document: dict[str, Any]) -> dict[str, Any]:
+        """Convert MongoDB document to memory format"""
+        if not document:
+            return {}
+
+        memory = document.copy()
+
+        # Convert ObjectId to string
+        if "_id" in memory:
+            memory["_id"] = str(memory["_id"])
+
+        # Convert datetime objects to ISO strings for JSON compatibility
+        datetime_fields = [
+            "created_at",
+            "expires_at",
+            "last_accessed",
+            "extraction_timestamp",
+            "timestamp",
+        ]
+        for field in datetime_fields:
+            if field in memory and isinstance(memory[field], datetime):
+                memory[field] = memory[field].isoformat()
+
+        return memory
+
+    # Chat History Operations
+    def store_chat_interaction(
+        self,
+        chat_id: str,
+        user_input: str,
+        ai_output: str,
+        model: str,
+        session_id: str,
+        namespace: str = "default",
+        tokens_used: int = 0,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
+        """Store a chat interaction in MongoDB"""
+        try:
+            collection = self.database[self.CHAT_HISTORY_COLLECTION]
+
+            document = {
+                "chat_id": chat_id,
+                "user_input": user_input,
+                "ai_output": ai_output,
+                "model": model,
+                "timestamp": datetime.now(timezone.utc),
+                "session_id": session_id,
+                "namespace": namespace,
+                "tokens_used": tokens_used,
+                "metadata": metadata or {},
+            }
+
+            collection.insert_one(document)
+            logger.debug(f"Stored chat interaction: {chat_id}")
+            return chat_id
+
+        except DuplicateKeyError:
+            # Chat ID already exists, update instead
+            return self.update_chat_interaction(
+                chat_id, user_input, ai_output, model, tokens_used, metadata
+            )
+        except Exception as e:
+            logger.error(f"Failed to store chat interaction: {e}")
+            raise
+
+    def update_chat_interaction(
+        self,
+        chat_id: str,
+        user_input: str,
+        ai_output: str,
+        model: str,
+        tokens_used: int = 0,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
+        """Update an existing chat interaction"""
+        try:
+            collection = self.database[self.CHAT_HISTORY_COLLECTION]
+
+            update_doc = {
+                "$set": {
+                    "user_input": user_input,
+                    "ai_output": ai_output,
+                    "model": model,
+                    "tokens_used": tokens_used,
+                    "timestamp": datetime.now(timezone.utc),
+                }
+            }
+
+            if metadata:
+                update_doc["$set"]["metadata"] = metadata
+
+            result = collection.update_one({"chat_id": chat_id}, update_doc)
+
+            if result.matched_count == 0:
+                raise ValueError(f"Chat interaction not found: {chat_id}")
+
+            logger.debug(f"Updated chat interaction: {chat_id}")
+            return chat_id
+
+        except Exception as e:
+            logger.error(f"Failed to update chat interaction: {e}")
+            raise
+
+    def get_chat_history(
+        self,
+        namespace: str = "default",
+        session_id: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """Retrieve chat history from MongoDB"""
+        try:
+            collection = self.database[self.CHAT_HISTORY_COLLECTION]
+
+            # Build filter
+            filter_doc = {"namespace": namespace}
+            if session_id:
+                filter_doc["session_id"] = session_id
+
+            # Execute query
+            cursor = collection.find(filter_doc).sort("timestamp", -1).limit(limit)
+
+            results = []
+            for document in cursor:
+                results.append(self._convert_document_to_memory(document))
+
+            logger.debug(f"Retrieved {len(results)} chat history entries")
+            return results
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve chat history: {e}")
+            return []
+
+    # Short-term Memory Operations
+    def store_short_term_memory(self, memory_data: dict[str, Any]) -> str:
+        """Store short-term memory in MongoDB"""
+        try:
+            collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+
+            # Generate memory ID if not provided
+            if "memory_id" not in memory_data:
+                memory_data["memory_id"] = str(uuid4())
+
+            document = self._convert_memory_to_document(memory_data)
+
+            collection.insert_one(document)
+            logger.debug(f"Stored short-term memory: {memory_data['memory_id']}")
+            return memory_data["memory_id"]
+
+        except Exception as e:
+            logger.error(f"Failed to store short-term memory: {e}")
+            raise
+
+    def get_short_term_memories(
+        self,
+        namespace: str = "default",
+        category_filter: list[str] | None = None,
+        importance_threshold: float = 0.0,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """Retrieve short-term memories from MongoDB"""
+        try:
+            collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+
+            # Build filter
+            filter_doc: dict[str, Any] = {
+                "namespace": namespace,
+                "importance_score": {"$gte": importance_threshold},
+            }
+
+            if category_filter:
+                filter_doc["category_primary"] = {"$in": category_filter}
+
+            # Include only non-expired memories
+            filter_doc["$or"] = [
+                {"expires_at": {"$exists": False}},
+                {"expires_at": None},
+                {"expires_at": {"$gt": datetime.now(timezone.utc)}},
+            ]
+
+            # Execute query
+            cursor = (
+                collection.find(filter_doc)
+                .sort([("importance_score", -1), ("created_at", -1)])
+                .limit(limit)
+            )
+
+            results = []
+            for document in cursor:
+                results.append(self._convert_document_to_memory(document))
+
+            logger.debug(f"Retrieved {len(results)} short-term memories")
+            return results
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve short-term memories: {e}")
+            return []
+
+    def update_short_term_memory(self, memory_id: str, updates: dict[str, Any]) -> bool:
+        """Update a short-term memory"""
+        try:
+            collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+
+            # Convert updates to document format
+            update_doc = {"$set": self._convert_memory_to_document(updates)}
+
+            result = collection.update_one({"memory_id": memory_id}, update_doc)
+
+            success = result.matched_count > 0
+            if success:
+                logger.debug(f"Updated short-term memory: {memory_id}")
+            else:
+                logger.warning(f"Short-term memory not found: {memory_id}")
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Failed to update short-term memory: {e}")
+            return False
+
+    def delete_short_term_memory(self, memory_id: str) -> bool:
+        """Delete a short-term memory"""
+        try:
+            collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+
+            result = collection.delete_one({"memory_id": memory_id})
+
+            success = result.deleted_count > 0
+            if success:
+                logger.debug(f"Deleted short-term memory: {memory_id}")
+            else:
+                logger.warning(f"Short-term memory not found: {memory_id}")
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Failed to delete short-term memory: {e}")
+            return False
+
+    # Long-term Memory Operations
+    def store_long_term_memory(self, memory_data: dict[str, Any]) -> str:
+        """Store long-term memory in MongoDB"""
+        try:
+            collection = self.database[self.LONG_TERM_MEMORY_COLLECTION]
+
+            # Generate memory ID if not provided
+            if "memory_id" not in memory_data:
+                memory_data["memory_id"] = str(uuid4())
+
+            document = self._convert_memory_to_document(memory_data)
+
+            collection.insert_one(document)
+            logger.debug(f"Stored long-term memory: {memory_data['memory_id']}")
+            return memory_data["memory_id"]
+
+        except Exception as e:
+            logger.error(f"Failed to store long-term memory: {e}")
+            raise
+
+    def get_long_term_memories(
+        self,
+        namespace: str = "default",
+        category_filter: list[str] | None = None,
+        importance_threshold: float = 0.0,
+        classification_filter: list[str] | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """Retrieve long-term memories from MongoDB"""
+        try:
+            collection = self.database[self.LONG_TERM_MEMORY_COLLECTION]
+
+            # Build filter
+            filter_doc = {
+                "namespace": namespace,
+                "importance_score": {"$gte": importance_threshold},
+            }
+
+            if category_filter:
+                filter_doc["category_primary"] = {"$in": category_filter}
+
+            if classification_filter:
+                filter_doc["classification"] = {"$in": classification_filter}
+
+            # Execute query
+            cursor = (
+                collection.find(filter_doc)
+                .sort([("importance_score", -1), ("created_at", -1)])
+                .limit(limit)
+            )
+
+            results = []
+            for document in cursor:
+                results.append(self._convert_document_to_memory(document))
+
+            logger.debug(f"Retrieved {len(results)} long-term memories")
+            return results
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve long-term memories: {e}")
+            return []
+
+    def update_long_term_memory(self, memory_id: str, updates: dict[str, Any]) -> bool:
+        """Update a long-term memory"""
+        try:
+            collection = self.database[self.LONG_TERM_MEMORY_COLLECTION]
+
+            # Convert updates to document format
+            update_doc = {"$set": self._convert_memory_to_document(updates)}
+
+            result = collection.update_one({"memory_id": memory_id}, update_doc)
+
+            success = result.matched_count > 0
+            if success:
+                logger.debug(f"Updated long-term memory: {memory_id}")
+            else:
+                logger.warning(f"Long-term memory not found: {memory_id}")
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Failed to update long-term memory: {e}")
+            return False
+
+    def delete_long_term_memory(self, memory_id: str) -> bool:
+        """Delete a long-term memory"""
+        try:
+            collection = self.database[self.LONG_TERM_MEMORY_COLLECTION]
+
+            result = collection.delete_one({"memory_id": memory_id})
+
+            success = result.deleted_count > 0
+            if success:
+                logger.debug(f"Deleted long-term memory: {memory_id}")
+            else:
+                logger.warning(f"Long-term memory not found: {memory_id}")
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Failed to delete long-term memory: {e}")
+            return False
+
+    # Search Operations
+    def search_memories(
+        self,
+        query: str,
+        namespace: str = "default",
+        memory_types: list[str] | None = None,
+        category_filter: list[str] | None = None,
+        limit: int = 10,
+    ) -> list[dict[str, Any]]:
+        """Search memories using MongoDB text search"""
+        try:
+            results = []
+            collections_to_search = []
+
+            # Determine which collections to search
+            if not memory_types or "short_term" in memory_types:
+                collections_to_search.append(
+                    (self.SHORT_TERM_MEMORY_COLLECTION, "short_term")
+                )
+            if not memory_types or "long_term" in memory_types:
+                collections_to_search.append(
+                    (self.LONG_TERM_MEMORY_COLLECTION, "long_term")
+                )
+
+            for collection_name, memory_type in collections_to_search:
+                collection = self.database[collection_name]
+
+                # Build search filter
+                search_filter: dict[str, Any] = {
+                    "$text": {"$search": query},
+                    "namespace": namespace,
+                }
+
+                if category_filter:
+                    search_filter["category_primary"] = {"$in": category_filter}
+
+                # For short-term memories, exclude expired ones
+                if memory_type == "short_term":
+                    search_filter["$or"] = [
+                        {"expires_at": {"$exists": False}},
+                        {"expires_at": None},
+                        {"expires_at": {"$gt": datetime.now(timezone.utc)}},
+                    ]
+
+                # Execute text search
+                cursor = (
+                    collection.find(search_filter, {"score": {"$meta": "textScore"}})
+                    .sort([("score", {"$meta": "textScore"}), ("importance_score", -1)])
+                    .limit(limit)
+                )
+
+                for document in cursor:
+                    memory = self._convert_document_to_memory(document)
+                    memory["memory_type"] = memory_type
+                    memory["search_strategy"] = "mongodb_text"
+                    results.append(memory)
+
+            # Sort all results by text score and importance
+            results.sort(
+                key=lambda x: (x.get("score", 0), x.get("importance_score", 0)),
+                reverse=True,
+            )
+
+            logger.debug(f"MongoDB text search returned {len(results)} results")
+            return results[:limit]
+
+        except Exception as e:
+            logger.error(f"MongoDB text search failed: {e}")
+            return self._fallback_search(query, namespace, category_filter, limit)
+
+    def _fallback_search(
+        self,
+        query: str,
+        namespace: str = "default",
+        category_filter: list[str] | None = None,
+        limit: int = 10,
+    ) -> list[dict[str, Any]]:
+        """Fallback search using regex when text search fails"""
+        try:
+            results = []
+            collections_to_search = [
+                (self.SHORT_TERM_MEMORY_COLLECTION, "short_term"),
+                (self.LONG_TERM_MEMORY_COLLECTION, "long_term"),
+            ]
+
+            # Create case-insensitive regex pattern
+            regex_pattern = {"$regex": query, "$options": "i"}
+
+            for collection_name, memory_type in collections_to_search:
+                collection = self.database[collection_name]
+
+                # Build search filter using regex
+                search_filter: dict[str, Any] = {
+                    "$or": [
+                        {"searchable_content": regex_pattern},
+                        {"summary": regex_pattern},
+                    ],
+                    "namespace": namespace,
+                }
+
+                if category_filter:
+                    search_filter["category_primary"] = {"$in": category_filter}
+
+                # For short-term memories, exclude expired ones
+                if memory_type == "short_term":
+                    search_filter["$and"] = [
+                        search_filter.get("$and", []),
+                        {
+                            "$or": [
+                                {"expires_at": {"$exists": False}},
+                                {"expires_at": None},
+                                {"expires_at": {"$gt": datetime.now(timezone.utc)}},
+                            ]
+                        },
+                    ]
+
+                # Execute regex search
+                cursor = (
+                    collection.find(search_filter)
+                    .sort([("importance_score", -1), ("created_at", -1)])
+                    .limit(limit)
+                )
+
+                for document in cursor:
+                    memory = self._convert_document_to_memory(document)
+                    memory["memory_type"] = memory_type
+                    memory["search_strategy"] = "regex_fallback"
+                    results.append(memory)
+
+            # Sort by importance score
+            results.sort(key=lambda x: x.get("importance_score", 0), reverse=True)
+
+            logger.debug(f"Regex fallback search returned {len(results)} results")
+            return results[:limit]
+
+        except Exception as e:
+            logger.error(f"Fallback search failed: {e}")
+            return []
+
+    # Batch Operations
+    def batch_store_memories(
+        self, memories: list[dict[str, Any]], memory_type: str = "short_term"
+    ) -> list[str]:
+        """Store multiple memories in batch"""
+        try:
+            if memory_type == "short_term":
+                collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+            elif memory_type == "long_term":
+                collection = self.database[self.LONG_TERM_MEMORY_COLLECTION]
+            else:
+                raise ValueError(f"Invalid memory type: {memory_type}")
+
+            # Prepare documents
+            documents = []
+            memory_ids = []
+
+            for memory_data in memories:
+                if "memory_id" not in memory_data:
+                    memory_data["memory_id"] = str(uuid4())
+
+                memory_ids.append(memory_data["memory_id"])
+                documents.append(self._convert_memory_to_document(memory_data))
+
+            # Insert all documents
+            result = collection.insert_many(documents, ordered=False)
+
+            logger.info(
+                f"Batch stored {len(result.inserted_ids)} {memory_type} memories"
+            )
+            return memory_ids
+
+        except Exception as e:
+            logger.error(f"Batch store failed: {e}")
+            return []
+
+    def cleanup_expired_memories(self, namespace: str = "default") -> int:
+        """Remove expired short-term memories"""
+        try:
+            collection = self.database[self.SHORT_TERM_MEMORY_COLLECTION]
+
+            # Delete expired memories
+            result = collection.delete_many(
+                {
+                    "namespace": namespace,
+                    "expires_at": {"$lt": datetime.now(timezone.utc)},
+                }
+            )
+
+            count = result.deleted_count
+            if count > 0:
+                logger.info(
+                    f"Cleaned up {count} expired memories from namespace: {namespace}"
+                )
+
+            return count
+
+        except Exception as e:
+            logger.error(f"Failed to cleanup expired memories: {e}")
+            return 0
+
+    def get_memory_stats(self, namespace: str = "default") -> dict[str, Any]:
+        """Get memory storage statistics"""
+        try:
+            stats = {
+                "namespace": namespace,
+                "short_term_count": 0,
+                "long_term_count": 0,
+                "chat_history_count": 0,
+                "total_size_bytes": 0,
+            }
+
+            # Count documents in each collection
+            stats["short_term_count"] = self.database[
+                self.SHORT_TERM_MEMORY_COLLECTION
+            ].count_documents({"namespace": namespace})
+
+            stats["long_term_count"] = self.database[
+                self.LONG_TERM_MEMORY_COLLECTION
+            ].count_documents({"namespace": namespace})
+
+            stats["chat_history_count"] = self.database[
+                self.CHAT_HISTORY_COLLECTION
+            ].count_documents({"namespace": namespace})
+
+            # Get database stats
+            db_stats = self.database.command("dbStats")
+            stats["total_size_bytes"] = db_stats.get("dataSize", 0)
+
+            return stats
+
+        except Exception as e:
+            logger.error(f"Failed to get memory stats: {e}")
+            return {"error": str(e)}