arc53--docsgpt/application/vectorstore/pgvector.py

import logging
from typing import List, Optional, Any, Dict

from psycopg.types.json import Jsonb

from application.core.settings import settings
from application.vectorstore.base import BaseVectorStore
from application.vectorstore.document_class import Document


class PGVectorStore(BaseVectorStore):
    def __init__(
        self,
        source_id: str = "",
        embeddings_key: str = "embeddings",
        table_name: str = "documents",
        decoded_token: Optional[str] = None,
        vector_column: str = "embedding",
        text_column: str = "text",
        metadata_column: str = "metadata",
        connection_string: str = None,
    ):
        super().__init__()
        # Store the source_id for use in add_chunk
        self._source_id = str(source_id).replace("application/indexes/", "").rstrip("/")
        self._embeddings_key = embeddings_key
        self._table_name = table_name
        self._vector_column = vector_column
        self._text_column = text_column
        self._metadata_column = metadata_column
        self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
        
        # Use provided connection string or fall back to settings.
        # If PGVECTOR_CONNECTION_STRING is not set but POSTGRES_URI is,
        # reuse the same cluster — normalize from SQLAlchemy dialect to libpq form.
        self._connection_string = connection_string or getattr(settings, 'PGVECTOR_CONNECTION_STRING', None)

        if not self._connection_string and getattr(settings, 'POSTGRES_URI', None):
            from application.core.db_uri import normalize_pgvector_connection_string
            self._connection_string = normalize_pgvector_connection_string(settings.POSTGRES_URI)

        if not self._connection_string:
            raise ValueError(
                "PostgreSQL connection string is required. "
                "Set PGVECTOR_CONNECTION_STRING or POSTGRES_URI in settings, "
                "or pass connection_string parameter."
            )

        try:
            import psycopg
            from pgvector.psycopg import register_vector
        except ImportError:
            raise ImportError(
                "Could not import required packages. "
                "Please install with `pip install 'psycopg[binary,pool]' pgvector`."
            )

        self._psycopg = psycopg
        self._register_vector = register_vector
        self._connection = None
        self._ensure_table_exists()

    def _get_connection(self):
        """Get or create database connection"""
        if self._connection is None or self._connection.closed:
            self._connection = self._psycopg.connect(self._connection_string)
            # Register pgvector types
            self._register_vector(self._connection)
        return self._connection

    def _ensure_table_exists(self):
        """Create table and enable pgvector extension if they don't exist"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        try:
            # Enable pgvector extension
            cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            
            embedding_dim = getattr(self._embedding, 'dimension', 768)
            
            # Create table with vector column
            create_table_query = f"""
            CREATE TABLE IF NOT EXISTS {self._table_name} (
                id SERIAL PRIMARY KEY,
                {self._text_column} TEXT NOT NULL,
                {self._vector_column} vector({embedding_dim}),
                {self._metadata_column} JSONB,
                source_id TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
            """
            cursor.execute(create_table_query)
            
            # Create index for vector similarity search
            index_query = f"""
            CREATE INDEX IF NOT EXISTS {self._table_name}_{self._vector_column}_idx 
            ON {self._table_name} USING ivfflat ({self._vector_column} vector_cosine_ops)
            WITH (lists = 100);
            """
            cursor.execute(index_query)
            
            # Create index for source_id filtering
            source_index_query = f"""
            CREATE INDEX IF NOT EXISTS {self._table_name}_source_id_idx
            ON {self._table_name} (source_id);
            """
            cursor.execute(source_index_query)

            # Functional GIN index backing keyword_search full-text queries.
            fts_index_query = f"""
            CREATE INDEX IF NOT EXISTS {self._table_name}_text_fts_idx
            ON {self._table_name} USING gin(to_tsvector('english', {self._text_column}));
            """
            cursor.execute(fts_index_query)

            conn.commit()
        except Exception as e:
            conn.rollback()
            logging.error(f"Error creating table: {e}")
            raise
        finally:
            cursor.close()

    def search(
        self,
        question: str,
        k: int = 2,
        *args,
        score_threshold: float = None,
        **kwargs,
    ) -> List[Document]:
        """Search for similar documents using vector similarity.

        Args:
            question: The query string.
            k: Maximum number of results.
            score_threshold: Optional cosine-similarity floor in ``[0, 1]``.
                Cosine distance = ``1 - similarity``; rows with similarity below
                the threshold (distance above ``1 - threshold``) are dropped.
        """
        query_vector = self._embedding.embed_query(question)

        conn = self._get_connection()
        cursor = conn.cursor()

        try:
            # Use cosine distance for similarity search with proper vector formatting
            search_query = f"""
            SELECT {self._text_column}, {self._metadata_column},
                   ({self._vector_column} <=> %s::vector) as distance
            FROM {self._table_name}
            WHERE source_id = %s
            ORDER BY {self._vector_column} <=> %s::vector
            LIMIT %s;
            """

            cursor.execute(search_query, (query_vector, self._source_id, query_vector, k))
            results = cursor.fetchall()

            max_distance = None
            if score_threshold is not None:
                max_distance = 1.0 - float(score_threshold)

            documents = []
            for text, metadata, distance in results:
                if max_distance is not None and distance is not None and distance > max_distance:
                    continue
                metadata = metadata or {}
                documents.append(Document(page_content=text, metadata=metadata))

            return documents
            
        except Exception as e:
            logging.error(f"Error searching documents: {e}", exc_info=True)
            return []
        finally:
            cursor.close()

    def keyword_search(self, question: str, k: int = 10) -> List[Document]:
        """Full-text keyword search using Postgres ``websearch_to_tsquery``.

        Returns the same ``Document`` shape as :meth:`search`. The question is
        bound as a query parameter (never interpolated) to prevent injection.
        """
        conn = self._get_connection()
        cursor = conn.cursor()

        try:
            keyword_query = f"""
            SELECT {self._text_column}, {self._metadata_column},
                   ts_rank(
                       to_tsvector('english', {self._text_column}),
                       websearch_to_tsquery('english', %s)
                   ) AS rank
            FROM {self._table_name}
            WHERE source_id = %s
              AND to_tsvector('english', {self._text_column})
                  @@ websearch_to_tsquery('english', %s)
            ORDER BY rank DESC
            LIMIT %s;
            """

            cursor.execute(
                keyword_query, (question, self._source_id, question, k)
            )
            results = cursor.fetchall()

            documents = []
            for text, metadata, _rank in results:
                metadata = metadata or {}
                documents.append(Document(page_content=text, metadata=metadata))

            return documents

        except Exception as e:
            logging.error(f"Error in keyword search: {e}", exc_info=True)
            return []
        finally:
            cursor.close()

    def add_texts(
        self,
        texts: List[str],
        metadatas: Optional[List[Dict[str, Any]]] = None,
        *args,
        **kwargs,
    ) -> List[str]:
        """Add texts with their embeddings to the vector store"""
        if not texts:
            return []

        embeddings = self._embedding.embed_documents(texts)
        metadatas = metadatas or [{}] * len(texts)
        
        conn = self._get_connection()
        cursor = conn.cursor()
        
        try:
            insert_query = f"""
            INSERT INTO {self._table_name} ({self._text_column}, {self._vector_column}, {self._metadata_column}, source_id)
            VALUES (%s, %s, %s, %s)
            RETURNING id;
            """
            
            inserted_ids = []
            for text, embedding, metadata in zip(texts, embeddings, metadatas):
                cursor.execute(
                    insert_query,
                    (text, embedding, Jsonb(metadata), self._source_id)
                )
                inserted_id = cursor.fetchone()[0]
                inserted_ids.append(str(inserted_id))
            
            conn.commit()
            return inserted_ids
            
        except Exception as e:
            conn.rollback()
            logging.error(f"Error adding texts: {e}")
            raise
        finally:
            cursor.close()

    def delete_index(self, *args, **kwargs):
        """Delete all documents for this source_id"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        try:
            delete_query = f"DELETE FROM {self._table_name} WHERE source_id = %s;"
            cursor.execute(delete_query, (self._source_id,))
            conn.commit()
            
        except Exception as e:
            conn.rollback()
            logging.error(f"Error deleting index: {e}")
            raise
        finally:
            cursor.close()

    def save_local(self, *args, **kwargs):
        """No-op for PostgreSQL - data is already persisted"""
        pass

    def get_chunks(self) -> List[Dict[str, Any]]:
        """Get all chunks for this source_id"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        try:
            select_query = f"""
            SELECT id, {self._text_column}, {self._metadata_column}
            FROM {self._table_name}
            WHERE source_id = %s;
            """
            cursor.execute(select_query, (self._source_id,))
            results = cursor.fetchall()
            
            chunks = []
            for doc_id, text, metadata in results:
                chunks.append({
                    "doc_id": str(doc_id),
                    "text": text,
                    "metadata": metadata or {}
                })
            
            return chunks
            
        except Exception as e:
            logging.error(f"Error getting chunks: {e}")
            return []
        finally:
            cursor.close()

    def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str:
        """Add a single chunk to the vector store"""
        metadata = metadata or {}

        final_metadata = metadata.copy()

        final_metadata["source_id"] = self._source_id

        embeddings = self._embedding.embed_documents([text])

        if not embeddings:
            raise ValueError("Could not generate embedding for chunk")
        
        conn = self._get_connection()
        cursor = conn.cursor()
        
        try:
            insert_query = f"""
            INSERT INTO {self._table_name} ({self._text_column}, {self._vector_column}, {self._metadata_column}, source_id)
            VALUES (%s, %s, %s, %s)
            RETURNING id;
            """
            
            cursor.execute(
                insert_query,
                (text, embeddings[0], Jsonb(final_metadata), self._source_id)
            )
            inserted_id = cursor.fetchone()[0]
            conn.commit()
            
            return str(inserted_id)
            
        except Exception as e:
            conn.rollback()
            logging.error(f"Error adding chunk: {e}")
            raise
        finally:
            cursor.close()

    def delete_chunk(self, chunk_id: str) -> bool:
        """Delete a specific chunk by its ID"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        try:
            delete_query = f"DELETE FROM {self._table_name} WHERE id = %s AND source_id = %s;"
            cursor.execute(delete_query, (int(chunk_id), self._source_id))
            deleted_count = cursor.rowcount
            conn.commit()
            
            return deleted_count > 0
            
        except Exception as e:
            conn.rollback()
            logging.error(f"Error deleting chunk: {e}")
            return False
        finally:
            cursor.close()

    def delete_chunks_by_source_path(self, path: str) -> int:
        """Delete this source's chunks whose ``metadata.source`` equals ``path``.

        One targeted statement instead of the base loop+scan. The path is bound
        as a query parameter (never interpolated); only the internal table name
        is f-string interpolated. Returns the number of rows deleted.
        """
        conn = self._get_connection()
        cursor = conn.cursor()

        try:
            delete_query = (
                f"DELETE FROM {self._table_name} "
                f"WHERE source_id = %s AND {self._metadata_column}->>'source' = %s;"
            )
            cursor.execute(delete_query, (self._source_id, path))
            deleted_count = cursor.rowcount
            conn.commit()

            return deleted_count

        except Exception as e:
            conn.rollback()
            logging.error(f"Error deleting chunks by source path: {e}")
            raise
        finally:
            cursor.close()

    def __del__(self):
        """Close database connection when object is destroyed"""
        if hasattr(self, '_connection') and self._connection and not self._connection.closed:
            self._connection.close()