mem0ai-azure-mysql 0.1.115__py3-none-any.whl

mem0/vector_stores/opensearch.py

@@ -0,0 +1,281 @@
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+try:
+    from opensearchpy import OpenSearch, RequestsHttpConnection
+except ImportError:
+    raise ImportError("OpenSearch requires extra dependencies. Install with `pip install opensearch-py`") from None
+
+from pydantic import BaseModel
+
+from mem0.configs.vector_stores.opensearch import OpenSearchConfig
+from mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: str
+    score: float
+    payload: Dict
+
+
+class OpenSearchDB(VectorStoreBase):
+    def __init__(self, **kwargs):
+        config = OpenSearchConfig(**kwargs)
+
+        # Initialize OpenSearch client
+        self.client = OpenSearch(
+            hosts=[{"host": config.host, "port": config.port or 9200}],
+            http_auth=config.http_auth
+            if config.http_auth
+            else ((config.user, config.password) if (config.user and config.password) else None),
+            use_ssl=config.use_ssl,
+            verify_certs=config.verify_certs,
+            connection_class=RequestsHttpConnection,
+            pool_maxsize=20,
+        )
+
+        self.collection_name = config.collection_name
+        self.embedding_model_dims = config.embedding_model_dims
+        self.create_col(self.collection_name, self.embedding_model_dims)
+
+    def create_index(self) -> None:
+        """Create OpenSearch index with proper mappings if it doesn't exist."""
+        index_settings = {
+            "settings": {
+                "index": {"number_of_replicas": 1, "number_of_shards": 5, "refresh_interval": "10s", "knn": True}
+            },
+            "mappings": {
+                "properties": {
+                    "text": {"type": "text"},
+                    "vector_field": {
+                        "type": "knn_vector",
+                        "dimension": self.embedding_model_dims,
+                        "method": {"engine": "nmslib", "name": "hnsw", "space_type": "cosinesimil"},
+                    },
+                    "metadata": {"type": "object", "properties": {"user_id": {"type": "keyword"}}},
+                }
+            },
+        }
+
+        if not self.client.indices.exists(index=self.collection_name):
+            self.client.indices.create(index=self.collection_name, body=index_settings)
+            logger.info(f"Created index {self.collection_name}")
+        else:
+            logger.info(f"Index {self.collection_name} already exists")
+
+    def create_col(self, name: str, vector_size: int) -> None:
+        """Create a new collection (index in OpenSearch)."""
+        index_settings = {
+            "settings": {"index.knn": True},
+            "mappings": {
+                "properties": {
+                    "vector_field": {
+                        "type": "knn_vector",
+                        "dimension": vector_size,
+                        "method": {"engine": "nmslib", "name": "hnsw", "space_type": "cosinesimil"},
+                    },
+                    "payload": {"type": "object"},
+                    "id": {"type": "keyword"},
+                }
+            },
+        }
+
+        if not self.client.indices.exists(index=name):
+            logger.warning(f"Creating index {name}, it might take 1-2 minutes...")
+            self.client.indices.create(index=name, body=index_settings)
+
+            # Wait for index to be ready
+            max_retries = 180  # 3 minutes timeout
+            retry_count = 0
+            while retry_count < max_retries:
+                try:
+                    # Check if index is ready by attempting a simple search
+                    self.client.search(index=name, body={"query": {"match_all": {}}})
+                    time.sleep(1)
+                    logger.info(f"Index {name} is ready")
+                    return
+                except Exception:
+                    retry_count += 1
+                    if retry_count == max_retries:
+                        raise TimeoutError(f"Index {name} creation timed out after {max_retries} seconds")
+                    time.sleep(0.5)
+
+    def insert(
+        self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None
+    ) -> List[OutputData]:
+        """Insert vectors into the index."""
+        if not ids:
+            ids = [str(i) for i in range(len(vectors))]
+
+        if payloads is None:
+            payloads = [{} for _ in range(len(vectors))]
+
+        for i, (vec, id_) in enumerate(zip(vectors, ids)):
+            body = {
+                "vector_field": vec,
+                "payload": payloads[i],
+                "id": id_,
+            }
+            self.client.index(index=self.collection_name, body=body)
+
+        results = []
+
+        return results
+
+    def search(
+        self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
+    ) -> List[OutputData]:
+        """Search for similar vectors using OpenSearch k-NN search with optional filters."""
+
+        # Base KNN query
+        knn_query = {
+            "knn": {
+                "vector_field": {
+                    "vector": vectors,
+                    "k": limit * 2,
+                }
+            }
+        }
+
+        # Start building the full query
+        query_body = {"size": limit * 2, "query": None}
+
+        # Prepare filter conditions if applicable
+        filter_clauses = []
+        if filters:
+            for key in ["user_id", "run_id", "agent_id"]:
+                value = filters.get(key)
+                if value:
+                    filter_clauses.append({"term": {f"payload.{key}.keyword": value}})
+
+        # Combine knn with filters if needed
+        if filter_clauses:
+            query_body["query"] = {"bool": {"must": knn_query, "filter": filter_clauses}}
+        else:
+            query_body["query"] = knn_query
+
+        # Execute search
+        response = self.client.search(index=self.collection_name, body=query_body)
+
+        hits = response["hits"]["hits"]
+        results = [
+            OutputData(id=hit["_source"].get("id"), score=hit["_score"], payload=hit["_source"].get("payload", {}))
+            for hit in hits
+        ]
+        return results
+
+    def delete(self, vector_id: str) -> None:
+        """Delete a vector by custom ID."""
+        # First, find the document by custom ID
+        search_query = {"query": {"term": {"id": vector_id}}}
+
+        response = self.client.search(index=self.collection_name, body=search_query)
+        hits = response.get("hits", {}).get("hits", [])
+
+        if not hits:
+            return
+
+        opensearch_id = hits[0]["_id"]
+
+        # Delete using the actual document ID
+        self.client.delete(index=self.collection_name, id=opensearch_id)
+
+    def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None) -> None:
+        """Update a vector and its payload using the custom 'id' field."""
+
+        # First, find the document by custom ID
+        search_query = {"query": {"term": {"id": vector_id}}}
+
+        response = self.client.search(index=self.collection_name, body=search_query)
+        hits = response.get("hits", {}).get("hits", [])
+
+        if not hits:
+            return
+
+        opensearch_id = hits[0]["_id"]  # The actual document ID in OpenSearch
+
+        # Prepare updated fields
+        doc = {}
+        if vector is not None:
+            doc["vector_field"] = vector
+        if payload is not None:
+            doc["payload"] = payload
+
+        if doc:
+            try:
+                response = self.client.update(index=self.collection_name, id=opensearch_id, body={"doc": doc})
+            except Exception:
+                pass
+
+    def get(self, vector_id: str) -> Optional[OutputData]:
+        """Retrieve a vector by ID."""
+        try:
+            # First check if index exists
+            if not self.client.indices.exists(index=self.collection_name):
+                logger.info(f"Index {self.collection_name} does not exist, creating it...")
+                self.create_col(self.collection_name, self.embedding_model_dims)
+                return None
+
+            search_query = {"query": {"term": {"id": vector_id}}}
+            response = self.client.search(index=self.collection_name, body=search_query)
+
+            hits = response["hits"]["hits"]
+
+            if not hits:
+                return None
+
+            return OutputData(id=hits[0]["_source"].get("id"), score=1.0, payload=hits[0]["_source"].get("payload", {}))
+        except Exception as e:
+            logger.error(f"Error retrieving vector {vector_id}: {str(e)}")
+            return None
+
+    def list_cols(self) -> List[str]:
+        """List all collections (indices)."""
+        return list(self.client.indices.get_alias().keys())
+
+    def delete_col(self) -> None:
+        """Delete a collection (index)."""
+        self.client.indices.delete(index=self.collection_name)
+
+    def col_info(self, name: str) -> Any:
+        """Get information about a collection (index)."""
+        return self.client.indices.get(index=name)
+
+    def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[OutputData]:
+        try:
+            """List all memories with optional filters."""
+            query: Dict = {"query": {"match_all": {}}}
+
+            filter_clauses = []
+            if filters:
+                for key in ["user_id", "run_id", "agent_id"]:
+                    value = filters.get(key)
+                    if value:
+                        filter_clauses.append({"term": {f"payload.{key}.keyword": value}})
+
+            if filter_clauses:
+                query["query"] = {"bool": {"filter": filter_clauses}}
+
+            if limit:
+                query["size"] = limit
+
+            response = self.client.search(index=self.collection_name, body=query)
+            hits = response["hits"]["hits"]
+
+            return [
+                [
+                    OutputData(id=hit["_source"].get("id"), score=1.0, payload=hit["_source"].get("payload", {}))
+                    for hit in hits
+                ]
+            ]
+        except Exception:
+            return []
+
+    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)

mem0/vector_stores/pgvector.py

@@ -0,0 +1,294 @@
+import json
+import logging
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+try:
+    import psycopg2
+    from psycopg2.extras import execute_values
+except ImportError:
+    raise ImportError("The 'psycopg2' library is required. Please install it using 'pip install psycopg2'.")
+
+from mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]
+    score: Optional[float]
+    payload: Optional[dict]
+
+
+class PGVector(VectorStoreBase):
+    def __init__(
+        self,
+        dbname,
+        collection_name,
+        embedding_model_dims,
+        user,
+        password,
+        host,
+        port,
+        diskann,
+        hnsw,
+    ):
+        """
+        Initialize the PGVector database.
+
+        Args:
+            dbname (str): Database name
+            collection_name (str): Collection name
+            embedding_model_dims (int): Dimension of the embedding vector
+            user (str): Database user
+            password (str): Database password
+            host (str, optional): Database host
+            port (int, optional): Database port
+            diskann (bool, optional): Use DiskANN for faster search
+            hnsw (bool, optional): Use HNSW for faster search
+        """
+        self.collection_name = collection_name
+        self.use_diskann = diskann
+        self.use_hnsw = hnsw
+        self.embedding_model_dims = embedding_model_dims
+
+        self.conn = psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port)
+        self.cur = self.conn.cursor()
+
+        collections = self.list_cols()
+        if collection_name not in collections:
+            self.create_col(embedding_model_dims)
+
+    def create_col(self, embedding_model_dims):
+        """
+        Create a new collection (table in PostgreSQL).
+        Will also initialize vector search index if specified.
+
+        Args:
+            embedding_model_dims (int): Dimension of the embedding vector.
+        """
+        self.cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+        self.cur.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.collection_name} (
+                id UUID PRIMARY KEY,
+                vector vector({embedding_model_dims}),
+                payload JSONB
+            );
+        """
+        )
+
+        if self.use_diskann and embedding_model_dims < 2000:
+            # Check if vectorscale extension is installed
+            self.cur.execute("SELECT * FROM pg_extension WHERE extname = 'vectorscale'")
+            if self.cur.fetchone():
+                # Create DiskANN index if extension is installed for faster search
+                self.cur.execute(
+                    f"""
+                    CREATE INDEX IF NOT EXISTS {self.collection_name}_diskann_idx
+                    ON {self.collection_name}
+                    USING diskann (vector);
+                """
+                )
+        elif self.use_hnsw:
+            self.cur.execute(
+                f"""
+                CREATE INDEX IF NOT EXISTS {self.collection_name}_hnsw_idx
+                ON {self.collection_name}
+                USING hnsw (vector vector_cosine_ops)
+            """
+            )
+
+        self.conn.commit()
+
+    def insert(self, vectors, payloads=None, ids=None):
+        """
+        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.
+        """
+        logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
+        json_payloads = [json.dumps(payload) for payload in payloads]
+
+        data = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, json_payloads)]
+        execute_values(
+            self.cur,
+            f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s",
+            data,
+        )
+        self.conn.commit()
+
+    def search(self, query, vectors, limit=5, filters=None):
+        """
+        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.
+        """
+        filter_conditions = []
+        filter_params = []
+
+        if filters:
+            for k, v in filters.items():
+                filter_conditions.append("payload->>%s = %s")
+                filter_params.extend([k, str(v)])
+
+        filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
+
+        self.cur.execute(
+            f"""
+            SELECT id, vector <=> %s::vector AS distance, payload
+            FROM {self.collection_name}
+            {filter_clause}
+            ORDER BY distance
+            LIMIT %s
+        """,
+            (vectors, *filter_params, limit),
+        )
+
+        results = self.cur.fetchall()
+        return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
+
+    def delete(self, vector_id):
+        """
+        Delete a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to delete.
+        """
+        self.cur.execute(f"DELETE FROM {self.collection_name} WHERE id = %s", (vector_id,))
+        self.conn.commit()
+
+    def update(self, vector_id, 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.
+        """
+        if vector:
+            self.cur.execute(
+                f"UPDATE {self.collection_name} SET vector = %s WHERE id = %s",
+                (vector, vector_id),
+            )
+        if payload:
+            self.cur.execute(
+                f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
+                (psycopg2.extras.Json(payload), vector_id),
+            )
+        self.conn.commit()
+
+    def get(self, vector_id) -> OutputData:
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to retrieve.
+
+        Returns:
+            OutputData: Retrieved vector.
+        """
+        self.cur.execute(
+            f"SELECT id, vector, payload FROM {self.collection_name} WHERE id = %s",
+            (vector_id,),
+        )
+        result = self.cur.fetchone()
+        if not result:
+            return None
+        return OutputData(id=str(result[0]), score=None, payload=result[2])
+
+    def list_cols(self) -> List[str]:
+        """
+        List all collections.
+
+        Returns:
+            List[str]: List of collection names.
+        """
+        self.cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
+        return [row[0] for row in self.cur.fetchall()]
+
+    def delete_col(self):
+        """Delete a collection."""
+        self.cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
+        self.conn.commit()
+
+    def col_info(self):
+        """
+        Get information about a collection.
+
+        Returns:
+            Dict[str, Any]: Collection information.
+        """
+        self.cur.execute(
+            f"""
+            SELECT 
+                table_name, 
+                (SELECT COUNT(*) FROM {self.collection_name}) as row_count,
+                (SELECT pg_size_pretty(pg_total_relation_size('{self.collection_name}'))) as total_size
+            FROM information_schema.tables 
+            WHERE table_schema = 'public' AND table_name = %s
+        """,
+            (self.collection_name,),
+        )
+        result = self.cur.fetchone()
+        return {"name": result[0], "count": result[1], "size": result[2]}
+
+    def list(self, filters=None, limit=100):
+        """
+        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.
+        """
+        filter_conditions = []
+        filter_params = []
+
+        if filters:
+            for k, v in filters.items():
+                filter_conditions.append("payload->>%s = %s")
+                filter_params.extend([k, str(v)])
+
+        filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
+
+        query = f"""
+            SELECT id, vector, payload
+            FROM {self.collection_name}
+            {filter_clause}
+            LIMIT %s
+        """
+
+        self.cur.execute(query, (*filter_params, limit))
+
+        results = self.cur.fetchall()
+        return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]]
+
+    def __del__(self):
+        """
+        Close the database connection when the object is deleted.
+        """
+        if hasattr(self, "cur"):
+            self.cur.close()
+        if hasattr(self, "conn"):
+            self.conn.close()
+
+    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.embedding_model_dims)