mem0ai--mem0/mem0/vector_stores/mongodb.py

import logging
from importlib.metadata import version
from typing import Any, Dict, List, Optional

from pydantic import BaseModel

try:
    from pymongo import MongoClient
    from pymongo.driver_info import DriverInfo
    from pymongo.errors import PyMongoError
    from pymongo.operations import SearchIndexModel
except ImportError:
    raise ImportError("The 'pymongo' library is required. Please install it using 'pip install pymongo'.")

from mem0.vector_stores.base import VectorStoreBase

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

_DRIVER_METADATA = DriverInfo(name="Mem0", version=version("mem0ai"))

class OutputData(BaseModel):
    id: Optional[str]
    score: Optional[float]
    payload: Optional[dict]


class MongoDB(VectorStoreBase):
    VECTOR_TYPE = "vector"
    SIMILARITY_METRIC = "cosine"

    def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
        """
        Initialize the MongoDB vector store with vector search capabilities.

        Args:
            db_name (str): Database name
            collection_name (str): Collection name
            embedding_model_dims (int): Dimension of the embedding vector
            mongo_uri (str): MongoDB connection URI
        """
        self.collection_name = collection_name
        self.embedding_model_dims = embedding_model_dims
        self.db_name = db_name

        self.client = MongoClient(mongo_uri, driver=_DRIVER_METADATA)
        self.db = self.client[db_name]
        self.collection = self.create_col()

    def create_col(self):
        """Create new collection with vector search index."""
        try:
            database = self.client[self.db_name]
            collection_names = database.list_collection_names(authorizedCollections=True)
            if self.collection_name not in collection_names:
                logger.info(f"Collection '{self.collection_name}' does not exist. Creating it now.")
                collection = database[self.collection_name]
                # Insert and remove a placeholder document to create the collection
                collection.insert_one({"_id": 0, "placeholder": True})
                collection.delete_one({"_id": 0})
                logger.info(f"Collection '{self.collection_name}' created successfully.")
            else:
                collection = database[self.collection_name]

            self.index_name = f"{self.collection_name}_vector_index"
            found_indexes = list(collection.list_search_indexes(name=self.index_name))
            if found_indexes:
                logger.info(f"Search index '{self.index_name}' already exists in collection '{self.collection_name}'.")
            else:
                search_index_model = SearchIndexModel(
                    name=self.index_name,
                    type="vectorSearch",
                    definition={
                        "fields": [
                            {
                                "type": self.VECTOR_TYPE,
                                "path": "embedding",
                                "numDimensions": self.embedding_model_dims,
                                "similarity": self.SIMILARITY_METRIC,
                            }
                        ]
                    },
                )
                collection.create_search_index(search_index_model)
                logger.info(
                    f"Search index '{self.index_name}' created successfully for collection '{self.collection_name}'."
                )

            # Create Atlas Search text index for keyword_search()
            text_index_name = f"{self.collection_name}_text_search_index"
            try:
                found_text_indexes = list(collection.list_search_indexes(name=text_index_name))
                if not found_text_indexes:
                    text_search_index_model = SearchIndexModel(
                        name=text_index_name,
                        definition={
                            "mappings": {
                                "dynamic": False,
                                "fields": {
                                    "payload": {
                                        "type": "document",
                                        "fields": {
                                            "data": {"type": "string"},
                                            "text_lemmatized": {"type": "string"},
                                        },
                                    }
                                },
                            }
                        },
                    )
                    collection.create_search_index(text_search_index_model)
                    logger.info(
                        f"Text search index '{text_index_name}' created successfully for collection '{self.collection_name}'."
                    )
                else:
                    logger.info(f"Text search index '{text_index_name}' already exists in collection '{self.collection_name}'.")
            except Exception as e:
                logger.warning(
                    f"Could not create text search index '{text_index_name}': {e}. "
                    "Atlas Search may not be available. keyword_search() will not work."
                )

            return collection
        except PyMongoError as e:
            logger.error(f"Error creating collection and search index: {e}")
            return None

    def insert(
        self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None
    ) -> None:
        """
        Insert vectors into the collection.

        Args:
            vectors (List[List[float]]): List of vectors to insert.
            payloads (List[Dict], optional): List of payloads corresponding to vectors.
            ids (List[str], optional): List of IDs corresponding to vectors.
        """
        logger.info(f"Inserting {len(vectors)} vectors into collection '{self.collection_name}'.")

        data = []
        for vector, payload, _id in zip(vectors, payloads or [{}] * len(vectors), ids or [None] * len(vectors)):
            document = {"_id": _id, "embedding": vector, "payload": payload}
            data.append(document)
        try:
            self.collection.insert_many(data)
            logger.info(f"Inserted {len(data)} documents into '{self.collection_name}'.")
        except PyMongoError as e:
            logger.error(f"Error inserting data: {e}")

    @staticmethod
    def _validate_filter_value(key: str, value: Any) -> None:
        """Reject values that could inject MongoDB query operators (e.g. $ne, $gt)."""
        if isinstance(value, dict):
            raise ValueError(
                f"Filter value for {key!r} must be a scalar (str, int, float, bool), "
                f"not a dict. Dicts may contain MongoDB query operators."
            )
        if isinstance(value, list):
            for item in value:
                if isinstance(item, dict):
                    raise ValueError(
                        f"Filter list for {key!r} contains a dict, "
                        f"which may contain MongoDB query operators."
                    )

    def search(self, query: str, vectors: List[float], top_k=5, filters: Optional[Dict] = None) -> List[OutputData]:
        """
        Search for similar vectors using the vector search index.

        Args:
            query (str): Query string
            vectors (List[float]): Query vector.
            top_k (int, optional): Number of results to return. Defaults to 5.
            filters (Dict, optional): Filters to apply to the search.

        Returns:
            List[OutputData]: Search results.
        """

        found_indexes = list(self.collection.list_search_indexes(name=self.index_name))
        if not found_indexes:
            logger.error(f"Index '{self.index_name}' does not exist.")
            return []

        if filters:
            for key, value in filters.items():
                self._validate_filter_value(key, value)

        results = []
        try:
            collection = self.client[self.db_name][self.collection_name]
            pipeline = [
                {
                    "$vectorSearch": {
                        "index": self.index_name,
                        "limit": top_k,
                        "numCandidates": min(top_k * 20, 10000),
                        "queryVector": vectors,
                        "path": "embedding",
                    }
                },
                {"$set": {"score": {"$meta": "vectorSearchScore"}}},
                {"$project": {"embedding": 0}},
            ]

            # Add filter stage if filters are provided
            if filters:
                filter_conditions = []
                for key, value in filters.items():
                    filter_conditions.append({"payload." + key: value})

                if filter_conditions:
                    # Add a $match stage after vector search to apply filters
                    pipeline.insert(1, {"$match": {"$and": filter_conditions}})

            results = list(collection.aggregate(pipeline))
            logger.info(f"Vector search completed. Found {len(results)} documents.")
        except Exception as e:
            logger.error(f"Error during vector search for query {query}: {e}")
            return []

        output = [OutputData(id=str(doc["_id"]), score=doc.get("score"), payload=doc.get("payload")) for doc in results]
        return output

    def keyword_search(self, query, top_k=5, filters=None):
        """
        Perform keyword-based search using MongoDB Atlas Search.

        Args:
            query (str): The text query to search for.
            top_k (int, optional): Number of results to return. Defaults to 5.
            filters (Dict, optional): Filters to apply to the search.

        Returns:
            List[OutputData]: Search results, or None if Atlas Search index is not available.
        """
        if filters:
            for key, value in filters.items():
                self._validate_filter_value(key, value)
        try:
            collection = self.client[self.db_name][self.collection_name]
            search_index_name = f"{self.collection_name}_text_search_index"

            pipeline = [
                {
                    "$search": {
                        "index": search_index_name,
                        "text": {
                            "query": query,
                            "path": ["payload.data", "payload.text_lemmatized"],
                        },
                    }
                },
                {"$set": {"score": {"$meta": "searchScore"}}},
                {"$project": {"embedding": 0}},
            ]

            # Add filter stage if filters are provided
            if filters:
                filter_conditions = []
                for key, value in filters.items():
                    filter_conditions.append({"payload." + key: value})
                if filter_conditions:
                    pipeline.insert(1, {"$match": {"$and": filter_conditions}})

            pipeline.append({"$limit": top_k})

            results = list(collection.aggregate(pipeline))
            logger.info(f"Keyword search completed. Found {len(results)} documents.")

            output = [
                OutputData(id=str(doc["_id"]), score=doc.get("score"), payload=doc.get("payload"))
                for doc in results
            ]
            return output
        except Exception as e:
            logger.error(f"Error during keyword search for query '{query}': {e}")
            return None

    def delete(self, vector_id: str) -> None:
        """
        Delete a vector by ID.

        Args:
            vector_id (str): ID of the vector to delete.
        """
        try:
            result = self.collection.delete_one({"_id": vector_id})
            if result.deleted_count > 0:
                logger.info(f"Deleted document with ID '{vector_id}'.")
            else:
                logger.warning(f"No document found with ID '{vector_id}' to delete.")
        except PyMongoError as e:
            logger.error(f"Error deleting document: {e}")

    def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None) -> None:
        """
        Update a vector and its payload.

        Args:
            vector_id (str): ID of the vector to update.
            vector (List[float], optional): Updated vector.
            payload (Dict, optional): Updated payload.
        """
        update_fields = {}
        if vector is not None:
            update_fields["embedding"] = vector
        if payload is not None:
            for key, value in payload.items():
                update_fields[f"payload.{key}"] = value

        if update_fields:
            try:
                result = self.collection.update_one({"_id": vector_id}, {"$set": update_fields})
                if result.matched_count > 0:
                    logger.info(f"Updated document with ID '{vector_id}'.")
                else:
                    logger.warning(f"No document found with ID '{vector_id}' to update.")
            except PyMongoError as e:
                logger.error(f"Error updating document: {e}")

    def get(self, vector_id: str) -> Optional[OutputData]:
        """
        Retrieve a vector by ID.

        Args:
            vector_id (str): ID of the vector to retrieve.

        Returns:
            Optional[OutputData]: Retrieved vector or None if not found.
        """
        try:
            doc = self.collection.find_one({"_id": vector_id})
            if doc:
                logger.info(f"Retrieved document with ID '{vector_id}'.")
                return OutputData(id=str(doc["_id"]), score=None, payload=doc.get("payload"))
            else:
                logger.warning(f"Document with ID '{vector_id}' not found.")
                return None
        except PyMongoError as e:
            logger.error(f"Error retrieving document: {e}")
            return None

    def list_cols(self) -> List[str]:
        """
        List all collections in the database.

        Returns:
            List[str]: List of collection names.
        """
        try:
            collections = self.db.list_collection_names(authorizedCollections=True)
            logger.info(f"Listing collections in database '{self.db_name}': {collections}")
            return collections
        except PyMongoError as e:
            logger.error(f"Error listing collections: {e}")
            return []

    def delete_col(self) -> None:
        """Delete the collection."""
        try:
            self.collection.drop()
            logger.info(f"Deleted collection '{self.collection_name}'.")
        except PyMongoError as e:
            logger.error(f"Error deleting collection: {e}")

    def col_info(self) -> Dict[str, Any]:
        """
        Get information about the collection.

        Returns:
            Dict[str, Any]: Collection information.
        """
        try:
            stats = self.db.command("collstats", self.collection_name)
            info = {"name": self.collection_name, "count": stats.get("count"), "size": stats.get("size")}
            logger.info(f"Collection info: {info}")
            return info
        except PyMongoError as e:
            logger.error(f"Error getting collection info: {e}")
            return {}

    def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
        """
        List vectors in the collection.

        Args:
            filters (Dict, optional): Filters to apply to the list.
            top_k (int, optional): Number of vectors to return.

        Returns:
            List[OutputData]: List of vectors.
        """
        if filters:
            for key, value in filters.items():
                self._validate_filter_value(key, value)
        try:
            query = {}
            if filters:
                # Apply filters to the payload field
                filter_conditions = []
                for key, value in filters.items():
                    filter_conditions.append({"payload." + key: value})
                if filter_conditions:
                    query = {"$and": filter_conditions}

            cursor = self.collection.find(query).limit(top_k)
            results = [OutputData(id=str(doc["_id"]), score=None, payload=doc.get("payload")) for doc in cursor]
            logger.info(f"Retrieved {len(results)} documents from collection '{self.collection_name}'.")
            return [results]
        except PyMongoError as e:
            logger.error(f"Error listing documents: {e}")
            return [[]]

    def reset(self):
        """Reset the index by deleting and recreating it."""
        logger.warning(f"Resetting index {self.collection_name}...")
        self.delete_col()
        self.collection = self.create_col()

    def __del__(self) -> None:
        """Close the database connection when the object is deleted."""
        if hasattr(self, "client"):
            self.client.close()
            logger.info("MongoClient connection closed.")