import os import json import psycopg2 from psycopg2.pool import ThreadedConnectionPool from abc import ABC, abstractmethod class VectorDBInterface(ABC): @abstractmethod def add_documents(self, ids, documents, metadatas, embeddings=None): pass @abstractmethod def query(self, query_embedding, n_results=2, where=None): pass @abstractmethod def get(self, ids): pass @abstractmethod def count(self): pass class ChromaDBClient(VectorDBInterface): def __init__(self, path, collection_name, embedding_function): import chromadb self.client = chromadb.PersistentClient(path=path) self.collection = self.client.get_or_create_collection( name=collection_name, embedding_function=embedding_function ) def add_documents(self, ids, documents, metadatas, embeddings=None): self.collection.add( ids=ids, documents=documents, metadatas=metadatas, embeddings=embeddings ) def query(self, query_embedding, n_results=2, where=None): query_kwargs = { "query_embeddings": [query_embedding], "n_results": n_results } if where: query_kwargs["where"] = where return self.collection.query(**query_kwargs) def get(self, ids): return self.collection.get(ids=ids) def count(self): return self.collection.count() class PGVectorClient(VectorDBInterface): def __init__(self, host, port, dbname, user, password, table_name="mintoak_content"): self.conn_params = { "host": host, "port": port, "dbname": dbname, "user": user, "password": password } self.table_name = table_name # Threaded connection pool: min 2, max 15 connections self.pool = ThreadedConnectionPool(2, 15, **self.conn_params) self._init_db() def _get_connection(self): return self.pool.getconn() def _release_connection(self, conn): self.pool.putconn(conn) def _init_db(self): conn = self._get_connection() try: with conn.cursor() as cur: cur.execute("CREATE EXTENSION IF NOT EXISTS vector;") cur.execute(f""" CREATE TABLE IF NOT EXISTS {self.table_name} ( id VARCHAR(255) PRIMARY KEY, document TEXT, metadata JSONB, embedding vector(384) ); """) # Create HNSW index for cosine distance to dramatically improve retrieval speed try: cur.execute(f""" CREATE INDEX IF NOT EXISTS {self.table_name}_hnsw_idx ON {self.table_name} USING hnsw (embedding vector_cosine_ops); """) except Exception as ie: print(f"Warning: could not create HNSW index: {ie}") conn.rollback() # Rollback failed index sub-transaction conn.commit() except Exception as e: print(f"Error initializing PGVector database: {e}") raise e finally: self._release_connection(conn) def _translate_where(self, where_dict): if not where_dict: return "", () conditions = [] params = [] for key, value in where_dict.items(): if isinstance(value, dict): for op, op_val in value.items(): if op == "$in": placeholders = ", ".join(["%s"] * len(op_val)) conditions.append(f"metadata->>%s IN ({placeholders})") params.append(key) params.extend(op_val) elif op == "$eq": conditions.append("metadata->>%s = %s") params.extend([key, op_val]) elif op == "$ne": conditions.append("metadata->>%s != %s") params.extend([key, op_val]) else: conditions.append("metadata->>%s = %s") params.extend([key, value]) if conditions: return " AND " + " AND ".join(conditions), tuple(params) return "", () def add_documents(self, ids, documents, metadatas, embeddings=None): if not embeddings: raise ValueError("Embeddings must be pre-computed and provided for pgvector insertion.") conn = self._get_connection() try: with conn.cursor() as cur: for idx, doc_id in enumerate(ids): doc = documents[idx] meta = metadatas[idx] emb = embeddings[idx] # Convert numpy array / list to string representation for vector type emb_list = emb.tolist() if hasattr(emb, "tolist") else list(emb) emb_str = "[" + ",".join(map(str, emb_list)) + "]" cur.execute(f""" INSERT INTO {self.table_name} (id, document, metadata, embedding) VALUES (%s, %s, %s, %s::vector) ON CONFLICT (id) DO UPDATE SET document = EXCLUDED.document, metadata = EXCLUDED.metadata, embedding = EXCLUDED.embedding; """, (doc_id, doc, json.dumps(meta), emb_str)) conn.commit() finally: self._release_connection(conn) def query(self, query_embedding, n_results=2, where=None): conn = self._get_connection() try: with conn.cursor() as cur: emb_str = "[" + ",".join(map(str, query_embedding)) + "]" where_clause, where_params = self._translate_where(where) query_str = f""" SELECT id, document, metadata, (embedding <=> %s::vector) AS distance FROM {self.table_name} WHERE 1=1 {where_clause} ORDER BY distance ASC LIMIT %s; """ params = (emb_str,) + where_params + (n_results,) cur.execute(query_str, params) rows = cur.fetchall() ids = [] documents = [] metadatas = [] distances = [] for row in rows: ids.append(row[0]) documents.append(row[1]) metadatas.append(row[2]) distances.append(row[3]) return { "ids": [ids], "documents": [documents], "metadatas": [metadatas], "distances": [distances] } finally: self._release_connection(conn) def get(self, ids): conn = self._get_connection() try: with conn.cursor() as cur: cur.execute(f""" SELECT id, document, metadata FROM {self.table_name} WHERE id = ANY(%s); """, (ids,)) rows = cur.fetchall() ret_ids = [] documents = [] metadatas = [] for row in rows: ret_ids.append(row[0]) documents.append(row[1]) metadatas.append(row[2]) return { "ids": ret_ids, "documents": documents, "metadatas": metadatas } finally: self._release_connection(conn) def count(self): conn = self._get_connection() try: with conn.cursor() as cur: cur.execute(f"SELECT COUNT(*) FROM {self.table_name};") return cur.fetchone()[0] finally: self._release_connection(conn) def get_vector_db(embedding_function=None): db_type = os.getenv("VECTOR_DB_TYPE", "chroma").lower() if db_type == "pgvector": host = os.getenv("PG_HOST", "localhost") port = os.getenv("PG_PORT", "5432") dbname = os.getenv("PG_NAME", "mintoak_db") user = os.getenv("PG_USER", "postgres") password = os.getenv("PG_PASSWORD", "postgres") return PGVectorClient(host, port, dbname, user, password) else: BASE_DIR = os.path.dirname(os.path.abspath(__file__)) CHROMA_DB_PATH = os.path.join(BASE_DIR, "data/mintoak/chroma_db") COLLECTION_NAME = "mintoak_content" return ChromaDBClient(CHROMA_DB_PATH, COLLECTION_NAME, embedding_function)