agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/vector_stores/milvus.py

@@ -0,0 +1,250 @@
+import logging
+from typing import Dict, Optional
+
+from pydantic import BaseModel
+
+from agentrun_mem0.configs.vector_stores.milvus import MetricType
+from agentrun_mem0.vector_stores.base import VectorStoreBase
+
+try:
+    import pymilvus  # noqa: F401
+except ImportError:
+    raise ImportError("The 'pymilvus' library is required. Please install it using 'pip install pymilvus'.")
+
+from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]  # memory id
+    score: Optional[float]  # distance
+    payload: Optional[Dict]  # metadata
+
+
+class MilvusDB(VectorStoreBase):
+    def __init__(
+        self,
+        url: str,
+        token: str,
+        collection_name: str,
+        embedding_model_dims: int,
+        metric_type: MetricType,
+        db_name: str,
+    ) -> None:
+        """Initialize the MilvusDB database.
+
+        Args:
+            url (str): Full URL for Milvus/Zilliz server.
+            token (str): Token/api_key for Zilliz server / for local setup defaults to None.
+            collection_name (str): Name of the collection (defaults to mem0).
+            embedding_model_dims (int): Dimensions of the embedding model (defaults to 1536).
+            metric_type (MetricType): Metric type for similarity search (defaults to L2).
+            db_name (str): Name of the database (defaults to "").
+        """
+        self.collection_name = collection_name
+        self.embedding_model_dims = embedding_model_dims
+        self.metric_type = metric_type
+        self.client = MilvusClient(uri=url, token=token, db_name=db_name)
+        self.create_col(
+            collection_name=self.collection_name,
+            vector_size=self.embedding_model_dims,
+            metric_type=self.metric_type,
+        )
+
+    def create_col(
+        self,
+        collection_name: str,
+        vector_size: int,
+        metric_type: MetricType = MetricType.COSINE,
+    ) -> None:
+        """Create a new collection with index_type AUTOINDEX.
+
+        Args:
+            collection_name (str): Name of the collection (defaults to mem0).
+            vector_size (int): Dimensions of the embedding model (defaults to 1536).
+            metric_type (MetricType, optional): etric type for similarity search. Defaults to MetricType.COSINE.
+        """
+
+        if self.client.has_collection(collection_name):
+            logger.info(f"Collection {collection_name} already exists. Skipping creation.")
+        else:
+            fields = [
+                FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512),
+                FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size),
+                FieldSchema(name="metadata", dtype=DataType.JSON),
+            ]
+
+            schema = CollectionSchema(fields, enable_dynamic_field=True)
+
+            index = self.client.prepare_index_params(
+                field_name="vectors", metric_type=metric_type, index_type="AUTOINDEX", index_name="vector_index"
+            )
+            self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index)
+
+    def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]):
+        """Insert vectors into a 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.
+        """
+        # Batch insert all records at once for better performance and consistency
+        data = [
+            {"id": idx, "vectors": embedding, "metadata": metadata}
+            for idx, embedding, metadata in zip(ids, vectors, payloads)
+        ]
+        self.client.insert(collection_name=self.collection_name, data=data, **kwargs)
+
+    def _create_filter(self, filters: dict):
+        """Prepare filters for efficient query.
+
+        Args:
+            filters (dict): filters [user_id, agent_id, run_id]
+
+        Returns:
+            str: formated filter.
+        """
+        operands = []
+        for key, value in filters.items():
+            if isinstance(value, str):
+                operands.append(f'(metadata["{key}"] == "{value}")')
+            else:
+                operands.append(f'(metadata["{key}"] == {value})')
+
+        return " and ".join(operands)
+
+    def _parse_output(self, data: list):
+        """
+        Parse the output data.
+
+        Args:
+            data (Dict): Output data.
+
+        Returns:
+            List[OutputData]: Parsed output data.
+        """
+        memory = []
+
+        for value in data:
+            uid, score, metadata = (
+                value.get("id"),
+                value.get("distance"),
+                value.get("entity", {}).get("metadata"),
+            )
+
+            memory_obj = OutputData(id=uid, score=score, payload=metadata)
+            memory.append(memory_obj)
+
+        return memory
+
+    def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
+        """
+        Search for similar vectors.
+
+        Args:
+            query (str): Query.
+            vectors (List[float]): Query vector.
+            limit (int, optional): Number of results to return. Defaults to 5.
+            filters (Dict, optional): Filters to apply to the search. Defaults to None.
+
+        Returns:
+            list: Search results.
+        """
+        query_filter = self._create_filter(filters) if filters else None
+        hits = self.client.search(
+            collection_name=self.collection_name,
+            data=[vectors],
+            limit=limit,
+            filter=query_filter,
+            output_fields=["*"],
+        )
+        result = self._parse_output(data=hits[0])
+        return result
+
+    def delete(self, vector_id):
+        """
+        Delete a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to delete.
+        """
+        self.client.delete(collection_name=self.collection_name, ids=vector_id)
+
+    def update(self, vector_id=None, vector=None, payload=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.
+        """
+        schema = {"id": vector_id, "vectors": vector, "metadata": payload}
+        self.client.upsert(collection_name=self.collection_name, data=schema)
+
+    def get(self, vector_id):
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to retrieve.
+
+        Returns:
+            OutputData: Retrieved vector.
+        """
+        result = self.client.get(collection_name=self.collection_name, ids=vector_id)
+        output = OutputData(
+            id=result[0].get("id", None),
+            score=None,
+            payload=result[0].get("metadata", None),
+        )
+        return output
+
+    def list_cols(self):
+        """
+        List all collections.
+
+        Returns:
+            List[str]: List of collection names.
+        """
+        return self.client.list_collections()
+
+    def delete_col(self):
+        """Delete a collection."""
+        return self.client.drop_collection(collection_name=self.collection_name)
+
+    def col_info(self):
+        """
+        Get information about a collection.
+
+        Returns:
+            Dict[str, Any]: Collection information.
+        """
+        return self.client.get_collection_stats(collection_name=self.collection_name)
+
+    def list(self, filters: dict = None, limit: int = 100) -> list:
+        """
+        List all vectors in a collection.
+
+        Args:
+            filters (Dict, optional): Filters to apply to the list.
+            limit (int, optional): Number of vectors to return. Defaults to 100.
+
+        Returns:
+            List[OutputData]: List of vectors.
+        """
+        query_filter = self._create_filter(filters) if filters else None
+        result = self.client.query(collection_name=self.collection_name, filter=query_filter, limit=limit)
+        memories = []
+        for data in result:
+            obj = OutputData(id=data.get("id"), score=None, payload=data.get("metadata"))
+            memories.append(obj)
+        return [memories]
+
+    def reset(self):
+        """Reset the index by deleting and recreating it."""
+        logger.warning(f"Resetting index {self.collection_name}...")
+        self.delete_col()
+        self.create_col(self.collection_name, self.embedding_model_dims, self.metric_type)

agentrun_mem0/vector_stores/mongodb.py

@@ -0,0 +1,310 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel
+
+try:
+    from pymongo import MongoClient
+    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 agentrun_mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]
+    score: Optional[float]
+    payload: Optional[dict]
+
+
+class MongoDB(VectorStoreBase):
+    VECTOR_TYPE = "knnVector"
+    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)
+        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()
+            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,
+                    definition={
+                        "mappings": {
+                            "dynamic": False,
+                            "fields": {
+                                "embedding": {
+                                    "type": self.VECTOR_TYPE,
+                                    "dimensions": 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}'."
+                )
+            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}")
+
+    def search(self, query: str, vectors: List[float], limit=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.
+            limit (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 []
+
+        results = []
+        try:
+            collection = self.client[self.db_name][self.collection_name]
+            pipeline = [
+                {
+                    "$vectorSearch": {
+                        "index": self.index_name,
+                        "limit": limit,
+                        "numCandidates": limit,
+                        "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 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:
+            update_fields["payload"] = payload
+
+        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()
+            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, limit: int = 100) -> List[OutputData]:
+        """
+        List vectors in the collection.
+
+        Args:
+            filters (Dict, optional): Filters to apply to the list.
+            limit (int, optional): Number of vectors to return.
+
+        Returns:
+            List[OutputData]: List of vectors.
+        """
+        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(limit)
+            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(self.collection_name)
+
+    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.")