Refactor PGVector to Use Internal Connection Pools and Context Managers (#3373)

mem0/configs/vector_stores/pgvector.py

@@ -11,19 +11,21 @@ class PGVectorConfig(BaseModel):
     password: Optional[str] = Field(None, description="Database password")
     host: Optional[str] = Field(None, description="Database host. Default is localhost")
     port: Optional[int] = Field(None, description="Database port. Default is 1536")
-    diskann: Optional[bool] = Field(True, description="Use diskann for approximate nearest neighbors search")
+    diskann: Optional[bool] = Field(False, description="Use diskann for approximate nearest neighbors search")
-    hnsw: Optional[bool] = Field(False, description="Use hnsw for faster search")
+    hnsw: Optional[bool] = Field(True, description="Use hnsw for faster search")
+    minconn: Optional[int] = Field(1, description="Minimum number of connections in the pool")
+    maxconn: Optional[int] = Field(5, description="Maximum number of connections in the pool")
     # New SSL and connection options
     sslmode: Optional[str] = Field(None, description="SSL mode for PostgreSQL connection (e.g., 'require', 'prefer', 'disable')")
     connection_string: Optional[str] = Field(None, description="PostgreSQL connection string (overrides individual connection parameters)")
-    connection_pool: Optional[Any] = Field(None, description="psycopg2 connection pool object (overrides connection string and individual parameters)")
+    connection_pool: Optional[Any] = Field(None, description="psycopg connection pool object (overrides connection string and individual parameters)")
 
     @model_validator(mode="before")
     def check_auth_and_connection(cls, values):
         # If connection_pool is provided, skip validation of individual connection parameters
         if values.get("connection_pool") is not None:
             return values
-        
+
         # If connection_string is provided, skip validation of individual connection parameters
         if values.get("connection_string") is not None:
             return values
@@ -32,9 +34,9 @@ class PGVectorConfig(BaseModel):
         user, password = values.get("user"), values.get("password")
         host, port = values.get("host"), values.get("port")
         if not user and not password:
-            raise ValueError("Both 'user' and 'password' must be provided when not using connection_string or connection_pool.")
+            raise ValueError("Both 'user' and 'password' must be provided when not using connection_string.")
         if not host and not port:
-            raise ValueError("Both 'host' and 'port' must be provided when not using connection_string or connection_pool.")
+            raise ValueError("Both 'host' and 'port' must be provided when not using connection_string.")
         return values
 
     @model_validator(mode="before")

mem0/vector_stores/pgvector.py

@@ -1,28 +1,28 @@
 import json
 import logging
-from typing import List, Optional
+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:
-    import psycopg
-    from psycopg import execute_values
     from psycopg.types.json import Json
+    from psycopg_pool import ConnectionPool
     PSYCOPG_VERSION = 3
     logger = logging.getLogger(__name__)
-    logger.info("Using psycopg (psycopg3) for PostgreSQL connections")
+    logger.info("Using psycopg (psycopg3) with ConnectionPool for PostgreSQL connections")
 except ImportError:
     try:
-        import psycopg2
+        from psycopg2.extras import Json, execute_values
-        from psycopg2.extras import execute_values, Json
+        from psycopg2.pool import ThreadedConnectionPool as ConnectionPool
         PSYCOPG_VERSION = 2
         logger = logging.getLogger(__name__)
-        logger.info("Using psycopg2 for PostgreSQL connections")
+        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' or 'pip install psycopg2'."
+            "Please install one of them using 'pip install psycopg[pool]' or 'pip install psycopg2'"
         )
 
 from mem0.vector_stores.base import VectorStoreBase
@@ -48,6 +48,8 @@ class PGVector(VectorStoreBase):
         port,
         diskann,
         hnsw,
+        minconn=1,
+        maxconn=5,
         sslmode=None,
         connection_string=None,
         connection_pool=None,
@@ -65,6 +67,8 @@ class PGVector(VectorStoreBase):
             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)
@@ -73,14 +77,13 @@ class PGVector(VectorStoreBase):
         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.conn = connection_pool.getconn()
             self.connection_pool = connection_pool
-        elif connection_string is not None:
+        elif connection_string:
-            # Use connection string
             if sslmode:
                 # Append sslmode to connection string if provided
                 if 'sslmode=' in connection_string:
@@ -90,99 +93,119 @@ class PGVector(VectorStoreBase):
                 else:
                     # Add sslmode to connection string
                     connection_string = f"{connection_string} sslmode={sslmode}"
-            
-            if PSYCOPG_VERSION == 3:
-                self.conn = psycopg.connect(connection_string)
-            else:
-                self.conn = psycopg2.connect(connection_string)
-            self.connection_pool = None
         else:
-            # Use individual connection parameters
+            connection_string = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
-            conn_params = {
-                'dbname': dbname,
-                'user': user,
-                'password': password,
-                'host': host,
-                'port': port
-            }
             if sslmode:
-                conn_params['sslmode'] = sslmode
+                connection_string = f"{connection_string} sslmode={sslmode}"
-            
-            if PSYCOPG_VERSION == 3:
-                self.conn = psycopg.connect(**conn_params)
-            else:
-                self.conn = psycopg2.connect(**conn_params)
-            self.connection_pool = None
         
-        self.cur = self.conn.cursor()
+        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(embedding_model_dims)
+            self.create_col()
 
-    def create_col(self, embedding_model_dims):
+    @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.
-
-        Args:
-            embedding_model_dims (int): Dimension of the embedding vector.
         """
-        self.cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+        with self._get_cursor(commit=True) as cur:
-        self.cur.execute(
+            cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
-            f"""
+            cur.execute(
-            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
+                CREATE TABLE IF NOT EXISTS {self.collection_name} (
-                ON {self.collection_name}
+                    id UUID PRIMARY KEY,
-                USING hnsw (vector vector_cosine_ops)
+                    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)
+                    """
+                )
 
-        self.conn.commit()
+    def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None:
-
-    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(
+        if PSYCOPG_VERSION == 3:
-            self.cur,
+            with self._get_cursor(commit=True) as cur:
-            f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s",
+                cur.executemany(
-            data,
+                    f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES (%s, %s, %s)",
-        )
+                    data,
-        self.conn.commit()
+                )
+        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, vectors, limit=5, filters=None):
+    def search(
+        self,
+        query: str,
+        vectors: list[float],
+        limit: Optional[int] = 5,
+        filters: Optional[dict] = None,
+    ) -> List[OutputData]:
         """
         Search for similar vectors.
 
@@ -205,31 +228,37 @@ class PGVector(VectorStoreBase):
 
         filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
 
-        self.cur.execute(
+        with self._get_cursor() as cur:
-            f"""
+            cur.execute(
-            SELECT id, vector <=> %s::vector AS distance, payload
+                f"""
-            FROM {self.collection_name}
+                SELECT id, vector <=> %s::vector AS distance, payload
-            {filter_clause}
+                FROM {self.collection_name}
-            ORDER BY distance
+                {filter_clause}
-            LIMIT %s
+                ORDER BY distance
-        """,
+                LIMIT %s
-            (vectors, *filter_params, limit),
+                """,
-        )
+                (vectors, *filter_params, limit),
+            )
 
-        results = self.cur.fetchall()
+            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):
+    def delete(self, vector_id: str) -> None:
         """
         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,))
+        with self._get_cursor(commit=True) as cur:
-        self.conn.commit()
+            cur.execute(f"DELETE FROM {self.collection_name} WHERE id = %s", (vector_id,))
 
-    def update(self, vector_id, vector=None, payload=None):
+    def update(
+        self,
+        vector_id: str,
+        vector: Optional[list[float]] = None,
+        payload: Optional[dict] = None,
+    ) -> None:
         """
         Update a vector and its payload.
 
@@ -238,28 +267,29 @@ class PGVector(VectorStoreBase):
             vector (List[float], optional): Updated vector.
             payload (Dict, optional): Updated payload.
         """
-        if vector:
+        with self._get_cursor(commit=True) as cur:
-            self.cur.execute(
+            if vector:
-                f"UPDATE {self.collection_name} SET vector = %s WHERE id = %s",
+               cur.execute(
-                (vector, vector_id),
+                    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
-                self.cur.execute(
-                    f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
-                    (Json(payload), vector_id),
                 )
-            else:
+            if payload:
-                # psycopg2 uses psycopg2.extras.Json
+                # Handle JSON serialization based on psycopg version
-                self.cur.execute(
+                if PSYCOPG_VERSION == 3:
-                    f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
+                    # psycopg3 uses psycopg.types.json.Json
-                    (psycopg2.extras.Json(payload), vector_id),
+                    cur.execute(
-                )
+                        f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
-        self.conn.commit()
+                        (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) -> OutputData:
+
+    def get(self, vector_id: str) -> OutputData:
         """
         Retrieve a vector by ID.
 
@@ -269,14 +299,15 @@ class PGVector(VectorStoreBase):
         Returns:
             OutputData: Retrieved vector.
         """
-        self.cur.execute(
+        with self._get_cursor() as cur:
-            f"SELECT id, vector, payload FROM {self.collection_name} WHERE id = %s",
+            cur.execute(
-            (vector_id,),
+                f"SELECT id, vector, payload FROM {self.collection_name} WHERE id = %s",
-        )
+                (vector_id,),
-        result = self.cur.fetchone()
+            )
-        if not result:
+            result = cur.fetchone()
-            return None
+            if not result:
-        return OutputData(id=str(result[0]), score=None, payload=result[2])
+                return None
+            return OutputData(id=str(result[0]), score=None, payload=result[2])
 
     def list_cols(self) -> List[str]:
         """
@@ -285,36 +316,42 @@ class PGVector(VectorStoreBase):
         Returns:
             List[str]: List of collection names.
         """
-        self.cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'")
+        with self._get_cursor() as cur:
-        return [row[0] for row in self.cur.fetchall()]
+            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):
+    def delete_col(self) -> None:
         """Delete a collection."""
-        self.cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
+        with self._get_cursor(commit=True) as cur:
-        self.conn.commit()
+            cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
 
-    def col_info(self):
+    def col_info(self) -> dict[str, Any]:
         """
         Get information about a collection.
 
         Returns:
             Dict[str, Any]: Collection information.
         """
-        self.cur.execute(
+        with self._get_cursor() as cur:
-            f"""
+            cur.execute(
-            SELECT 
+                f"""
-                table_name, 
+                SELECT
-                (SELECT COUNT(*) FROM {self.collection_name}) as row_count,
+                    table_name,
-                (SELECT pg_size_pretty(pg_total_relation_size('{self.collection_name}'))) as total_size
+                    (SELECT COUNT(*) FROM {self.collection_name}) as row_count,
-            FROM information_schema.tables 
+                    (SELECT pg_size_pretty(pg_total_relation_size('{self.collection_name}'))) as total_size
-            WHERE table_schema = 'public' AND table_name = %s
+                FROM information_schema.tables
-        """,
+                WHERE table_schema = 'public' AND table_name = %s
-            (self.collection_name,),
+            """,
-        )
+                (self.collection_name,),
-        result = self.cur.fetchone()
+            )
+            result = cur.fetchone()
         return {"name": result[0], "count": result[1], "size": result[2]}
 
-    def list(self, filters=None, limit=100):
+    def list(
+        self,
+        filters: Optional[dict] = None,
+        limit: Optional[int] = 100
+    ) -> List[OutputData]:
         """
         List all vectors in a collection.
 
@@ -342,27 +379,26 @@ class PGVector(VectorStoreBase):
             LIMIT %s
         """
 
-        self.cur.execute(query, (*filter_params, limit))
+        with self._get_cursor() as cur:
-
+            cur.execute(query, (*filter_params, limit))
-        results = self.cur.fetchall()
+            results = cur.fetchall()
         return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]]
 
-    def __del__(self):
+    def __del__(self) -> None:
         """
-        Close the database connection when the object is deleted.
+        Close the database connection pool when the object is deleted.
         """
-        if hasattr(self, "cur"):
+        try:
-            self.cur.close()
+            # Close pool appropriately
-        if hasattr(self, "conn"):
+            if PSYCOPG_VERSION == 3:
-            if hasattr(self, "connection_pool") and self.connection_pool is not None:
+                self.connection_pool.close()
-                # Return connection to pool instead of closing it
-                self.connection_pool.putconn(self.conn)
             else:
-                # Close the connection directly
+                self.connection_pool.closeall()
-                self.conn.close()
+        except Exception:
+            pass
 
-    def reset(self):
+    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(self.embedding_model_dims)
+        self.create_col()

pyproject.toml

@@ -39,6 +39,7 @@ vector_stores = [
     "upstash-vector>=0.1.0",
     "azure-search-documents>=11.4.0b8",
     "psycopg>=3.2.8",
+    "psycopg-pool>=3.2.6,<4.0.0",
     "pymongo>=4.13.2",
     "pymochow>=2.2.9",
     "databricks-sdk>=0.63.0",