agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/vector_stores/opensearch.py

@@ -0,0 +1,292 @@
+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 agentrun_mem0.configs.vector_stores.opensearch import OpenSearchConfig
+from agentrun_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))]
+
+        results = []
+        for i, (vec, id_) in enumerate(zip(vectors, ids)):
+            body = {
+                "vector_field": vec,
+                "payload": payloads[i],
+                "id": id_,
+            }
+            try:
+                self.client.index(index=self.collection_name, body=body)
+                # Force refresh to make documents immediately searchable for tests
+                self.client.indices.refresh(index=self.collection_name)
+                
+                results.append(OutputData(
+                    id=id_,
+                    score=1.0,  # No score for inserts
+                    payload=payloads[i]
+                ))
+            except Exception as e:
+                logger.error(f"Error inserting vector {id_}: {e}")
+                raise
+
+        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
+
+        try:
+            # 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[:limit]  # Ensure we don't exceed limit
+            ]
+            return results
+        except Exception as e:
+            logger.error(f"Error during search: {e}")
+            return []
+
+    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:
+            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 a flat list, not a nested array
+            results = [
+                OutputData(id=hit["_source"].get("id"), score=1.0, payload=hit["_source"].get("payload", {}))
+                for hit in hits
+            ]
+            return [results]  # VectorStore expects tuple/list format
+        except Exception as e:
+            logger.error(f"Error listing vectors: {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.create_col(self.collection_name, self.embedding_model_dims)

agentrun_mem0/vector_stores/pgvector.py

@@ -0,0 +1,404 @@
+import json
+import logging
+from contextlib import contextmanager
+from typing import Any, List, Optional
+
+from pydantic import BaseModel
+
+# Try to import psycopg (psycopg3) first, then fall back to psycopg2
+try:
+    from psycopg.types.json import Json
+    from psycopg_pool import ConnectionPool
+    PSYCOPG_VERSION = 3
+    logger = logging.getLogger(__name__)
+    logger.info("Using psycopg (psycopg3) with ConnectionPool for PostgreSQL connections")
+except ImportError:
+    try:
+        from psycopg2.extras import Json, execute_values
+        from psycopg2.pool import ThreadedConnectionPool as ConnectionPool
+        PSYCOPG_VERSION = 2
+        logger = logging.getLogger(__name__)
+        logger.info("Using psycopg2 with ThreadedConnectionPool for PostgreSQL connections")
+    except ImportError:
+        raise ImportError(
+            "Neither 'psycopg' nor 'psycopg2' library is available. "
+            "Please install one of them using 'pip install psycopg[pool]' or 'pip install psycopg2'"
+        )
+
+from agentrun_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,
+        minconn=1,
+        maxconn=5,
+        sslmode=None,
+        connection_string=None,
+        connection_pool=None,
+    ):
+        """
+        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
+            minconn (int): Minimum number of connections to keep in the connection pool
+            maxconn (int): Maximum number of connections allowed in the connection pool
+            sslmode (str, optional): SSL mode for PostgreSQL connection (e.g., 'require', 'prefer', 'disable')
+            connection_string (str, optional): PostgreSQL connection string (overrides individual connection parameters)
+            connection_pool (Any, optional): psycopg2 connection pool object (overrides connection string and individual parameters)
+        """
+        self.collection_name = collection_name
+        self.use_diskann = diskann
+        self.use_hnsw = hnsw
+        self.embedding_model_dims = embedding_model_dims
+        self.connection_pool = None
+
+        # Connection setup with priority: connection_pool > connection_string > individual parameters
+        if connection_pool is not None:
+            # Use provided connection pool
+            self.connection_pool = connection_pool
+        elif connection_string:
+            if sslmode:
+                # Append sslmode to connection string if provided
+                if 'sslmode=' in connection_string:
+                    # Replace existing sslmode
+                    import re
+                    connection_string = re.sub(r'sslmode=[^ ]*', f'sslmode={sslmode}', connection_string)
+                else:
+                    # Add sslmode to connection string
+                    connection_string = f"{connection_string} sslmode={sslmode}"
+        else:
+            connection_string = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
+            if sslmode:
+                connection_string = f"{connection_string} sslmode={sslmode}"
+        
+        if self.connection_pool is None:
+            if PSYCOPG_VERSION == 3:
+                # psycopg3 ConnectionPool
+                self.connection_pool = ConnectionPool(conninfo=connection_string, min_size=minconn, max_size=maxconn, open=True)
+            else:
+                # psycopg2 ThreadedConnectionPool
+                self.connection_pool = ConnectionPool(minconn=minconn, maxconn=maxconn, dsn=connection_string)
+
+        collections = self.list_cols()
+        if collection_name not in collections:
+            self.create_col()
+
+    @contextmanager
+    def _get_cursor(self, commit: bool = False):
+        """
+        Unified context manager to get a cursor from the appropriate pool.
+        Auto-commits or rolls back based on exception, and returns the connection to the pool.
+        """
+        if PSYCOPG_VERSION == 3:
+            # psycopg3 auto-manages commit/rollback and pool return
+            with self.connection_pool.connection() as conn:
+                with conn.cursor() as cur:
+                    try:
+                        yield cur
+                        if commit:
+                            conn.commit()
+                    except Exception:
+                        conn.rollback()
+                        logger.error("Error in cursor context (psycopg3)", exc_info=True)
+                        raise
+        else:
+            # psycopg2 manual getconn/putconn
+            conn = self.connection_pool.getconn()
+            cur = conn.cursor()
+            try:
+                yield cur
+                if commit:
+                    conn.commit()
+            except Exception as exc:
+                conn.rollback()
+                logger.error(f"Error occurred: {exc}")
+                raise exc
+            finally:
+                cur.close()
+                self.connection_pool.putconn(conn)
+
+    def create_col(self) -> None:
+        """
+        Create a new collection (table in PostgreSQL).
+        Will also initialize vector search index if specified.
+        """
+        with self._get_cursor(commit=True) as cur:
+            cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+            cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {self.collection_name} (
+                    id UUID PRIMARY KEY,
+                    vector vector({self.embedding_model_dims}),
+                    payload JSONB
+                );
+                """
+            )
+            if self.use_diskann and self.embedding_model_dims < 2000:
+                cur.execute("SELECT * FROM pg_extension WHERE extname = 'vectorscale'")
+                if cur.fetchone():
+                    # Create DiskANN index if extension is installed for faster search
+                    cur.execute(
+                        f"""
+                        CREATE INDEX IF NOT EXISTS {self.collection_name}_diskann_idx
+                        ON {self.collection_name}
+                        USING diskann (vector);
+                        """
+                    )
+            elif self.use_hnsw:
+                cur.execute(
+                    f"""
+                    CREATE INDEX IF NOT EXISTS {self.collection_name}_hnsw_idx
+                    ON {self.collection_name}
+                    USING hnsw (vector vector_cosine_ops)
+                    """
+                )
+
+    def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None:
+        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)]
+        if PSYCOPG_VERSION == 3:
+            with self._get_cursor(commit=True) as cur:
+                cur.executemany(
+                    f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES (%s, %s, %s)",
+                    data,
+                )
+        else:
+            with self._get_cursor(commit=True) as cur:
+                execute_values(
+                    cur,
+                    f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s",
+                    data,
+                )
+
+    def search(
+        self,
+        query: str,
+        vectors: list[float],
+        limit: Optional[int] = 5,
+        filters: Optional[dict] = None,
+    ) -> List[OutputData]:
+        """
+        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 ""
+
+        with self._get_cursor() as cur:
+            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 = cur.fetchall()
+        return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
+
+    def delete(self, vector_id: str) -> None:
+        """
+        Delete a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to delete.
+        """
+        with self._get_cursor(commit=True) as cur:
+            cur.execute(f"DELETE FROM {self.collection_name} WHERE id = %s", (vector_id,))
+
+    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.
+        """
+        with self._get_cursor(commit=True) as cur:
+            if vector:
+               cur.execute(
+                    f"UPDATE {self.collection_name} SET vector = %s WHERE id = %s",
+                    (vector, vector_id),
+                )
+            if payload:
+                # Handle JSON serialization based on psycopg version
+                if PSYCOPG_VERSION == 3:
+                    # psycopg3 uses psycopg.types.json.Json
+                    cur.execute(
+                        f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
+                        (Json(payload), vector_id),
+                    )
+                else:
+                    # psycopg2 uses psycopg2.extras.Json
+                    cur.execute(
+                        f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
+                        (Json(payload), vector_id),
+                    )
+
+
+    def get(self, vector_id: str) -> OutputData:
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to retrieve.
+
+        Returns:
+            OutputData: Retrieved vector.
+        """
+        with self._get_cursor() as cur:
+            cur.execute(
+                f"SELECT id, vector, payload FROM {self.collection_name} WHERE id = %s",
+                (vector_id,),
+            )
+            result = 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.
+        """
+        with self._get_cursor() as cur:
+            cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
+            return [row[0] for row in cur.fetchall()]
+
+    def delete_col(self) -> None:
+        """Delete a collection."""
+        with self._get_cursor(commit=True) as cur:
+            cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
+
+    def col_info(self) -> dict[str, Any]:
+        """
+        Get information about a collection.
+
+        Returns:
+            Dict[str, Any]: Collection information.
+        """
+        with self._get_cursor() as cur:
+            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 = cur.fetchone()
+        return {"name": result[0], "count": result[1], "size": result[2]}
+
+    def list(
+        self,
+        filters: Optional[dict] = None,
+        limit: Optional[int] = 100
+    ) -> List[OutputData]:
+        """
+        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
+        """
+
+        with self._get_cursor() as cur:
+            cur.execute(query, (*filter_params, limit))
+            results = cur.fetchall()
+        return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]]
+
+    def __del__(self) -> None:
+        """
+        Close the database connection pool when the object is deleted.
+        """
+        try:
+            # Close pool appropriately
+            if PSYCOPG_VERSION == 3:
+                self.connection_pool.close()
+            else:
+                self.connection_pool.closeall()
+        except Exception:
+            pass
+
+    def reset(self) -> None:
+        """Reset the index by deleting and recreating it."""
+        logger.warning(f"Resetting index {self.collection_name}...")
+        self.delete_col()
+        self.create_col()