vector_db.py

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)