JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Sheharyar Ahmad
2025-08-29 18:36:39 +05:00
committed by GitHub
parent 458d7ab8a3
commit 3d0ece1bcf
4 changed files with 1454 additions and 389 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/configs/vector_stores/pgvector.py
+8 -6
View File
@@ -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")
hnsw: Optional[bool] = Field(False, description="Use hnsw for faster search")
diskann: Optional[bool] = Field(False, description="Use diskann for approximate nearest neighbors 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
+198 -162
View File
@@ -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 execute_values, Json
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 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:
# Use connection string
elif 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
conn_params = {
'dbname': dbname,
'user': user,
'password': password,
'host': host,
'port': port
}
connection_string = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
if sslmode:
conn_params['sslmode'] = sslmode
if PSYCOPG_VERSION == 3:
self.conn = psycopg.connect(**conn_params)
else:
self.conn = psycopg2.connect(**conn_params)
self.connection_pool = None
connection_string = f"{connection_string} sslmode={sslmode}"
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")
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(
with self._get_cursor(commit=True) as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
cur.execute(
f"""
CREATE INDEX IF NOT EXISTS {self.collection_name}_hnsw_idx
ON {self.collection_name}
USING hnsw (vector vector_cosine_ops)
"""
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)
"""
)
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.
"""
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)]
execute_values(
self.cur,
f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s",
data,
)
self.conn.commit()
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, 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(
f"""
SELECT id, vector <=> %s::vector AS distance, payload
FROM {self.collection_name}
{filter_clause}
ORDER BY distance
LIMIT %s
""",
(vectors, *filter_params, limit),
)
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 = 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,))
self.conn.commit()
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, 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:
self.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
self.cur.execute(
f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
(Json(payload), vector_id),
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),
)
else:
# psycopg2 uses psycopg2.extras.Json
self.cur.execute(
f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
(psycopg2.extras.Json(payload), vector_id),
)
self.conn.commit()
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) -> 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(
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])
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]:
"""
@@ -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'")
return [row[0] for row in self.cur.fetchall()]
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):
def delete_col(self) -> None:
"""Delete a collection."""
self.cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
self.conn.commit()
with self._get_cursor(commit=True) as cur:
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(
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()
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=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))
results = self.cur.fetchall()
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):
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"):
self.cur.close()
if hasattr(self, "conn"):
if hasattr(self, "connection_pool") and self.connection_pool is not None:
# Return connection to pool instead of closing it
self.connection_pool.putconn(self.conn)
try:
# Close pool appropriately
if PSYCOPG_VERSION == 3:
self.connection_pool.close()
else:
# Close the connection directly
self.conn.close()
self.connection_pool.closeall()
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
+1
View File
@@ -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",
tests/vector_stores/test_pgvector.py
+1247 -221
View File
File diff suppressed because it is too large Load Diff