DocsGPT/application/vectorstore/pgvector.py

import logging
import math
import re
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

# table name -> IVFFlat ``lists`` (None when the table has no such index)
_IVFFLAT_LISTS_CACHE: Dict[str, Optional[int]] = {}


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)
            self._apply_ivfflat_probes(self._connection)
        return self._connection

    def _ivfflat_lists(self, conn) -> Optional[int]:
        """Return the ``lists`` value of this table's IVFFlat index, if any.

        Cached per table because it only changes when the index is rebuilt.
        """
        if self._table_name in _IVFFLAT_LISTS_CACHE:
            return _IVFFLAT_LISTS_CACHE[self._table_name]
        lists = None
        try:
            with conn.cursor() as cursor:
                cursor.execute(
                    "SELECT indexdef FROM pg_indexes "
                    "WHERE tablename = %s AND indexdef ILIKE %s",
                    (self._table_name, "%ivfflat%"),
                )
                row = cursor.fetchone()
            if row:
                match = re.search(r"lists\s*=\s*'?(\d+)", row[0])
                if match:
                    lists = int(match.group(1))
        except Exception as e:  # index introspection must never break search
            logging.debug("Could not read IVFFlat lists for %s: %s", self._table_name, e)
        # Only cache a hit: an index may be created after this process booted,
        # and caching None would keep probes unset for the process's lifetime.
        if lists:
            _IVFFLAT_LISTS_CACHE[self._table_name] = lists
        return lists

    def _apply_ivfflat_probes(self, conn) -> None:
        """Raise ``ivfflat.probes`` so a filtered search cannot come back empty.

        An IVFFlat index partitions vectors into ``lists`` clusters and the
        default ``probes = 1`` scans exactly one of them. Our searches filter by
        ``source_id`` *after* the index picks candidates, so with one probe the
        candidates frequently all belong to other sources and the query returns
        nothing — retrieval reports zero documents and the model answers with no
        source material, silently. ``sqrt(lists)`` is pgvector's own recall
        guidance and costs a proportional amount of scan.
        """
        probes = settings.PGVECTOR_IVFFLAT_PROBES
        if probes is None:
            lists = self._ivfflat_lists(conn)
            if not lists:
                return
            probes = max(1, math.isqrt(lists))
        try:
            with conn.cursor() as cursor:
                cursor.execute(f"SET ivfflat.probes = {int(probes)};")
        except Exception as e:  # older pgvector / no index — search still works
            logging.debug("Could not set ivfflat.probes: %s", e)

    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)
            
            # NO vector index is created here, deliberately.
            #
            # This runs when the table is first created, i.e. on an EMPTY
            # table. IVFFlat computes its cluster centroids at build time, so
            # an index built on no rows gets random centroids and never
            # recovers — measured recall 0.06 once 5k rows are added. Combined
            # with our ``WHERE source_id = ...`` post-filter, that returned
            # ZERO rows for sources with hundreds of chunks: retrieval reported
            # no documents and the model answered from nothing, silently.
            # pgvector only warns when sampled_rows < lists, so the common bad
            # case is silent.
            #
            # Exact search is correct and fast well past the sizes most
            # deployments ever reach. Add an index deliberately, sized to real
            # data, once a corpus is large enough to need one.
            # 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()

    score_kind = "cosine_similarity"

    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.
        """
        return [
            doc
            for doc, _ in self.search_with_scores(
                question, k, *args, score_threshold=score_threshold, **kwargs
            )
        ]

    def _nearest_sql(self) -> str:
        """Build the nearest-neighbour SELECT for this store's table.

        Identifiers (table/column names) come from this instance's
        construction, never from a request, so they cannot be interpolated by a
        caller; the query *values* are always bound parameters.
        """
        return (
            f"SELECT {self._text_column}, {self._metadata_column}, "
            f"({self._vector_column} <=> %s::vector) AS distance "
            f"FROM {self._table_name} "
            "WHERE source_id = %s "
            f"ORDER BY {self._vector_column} <=> %s::vector "
            "LIMIT %s;"
        )

    def _exact_search(self, cursor, query_vector, k: int, ann_results: list) -> list:
        """Redo a short indexed search exactly, when the source has more rows.

        Args:
            cursor: Open cursor on the search connection.
            query_vector: The embedded query.
            k: Requested top-k.
            ann_results: What the indexed search returned.

        Returns:
            list: Exact rows when the indexed search under-returned, otherwise
            ``ann_results`` unchanged.
        """
        try:
            cursor.execute(
                f"SELECT count(*) FROM {self._table_name} WHERE source_id = %s",
                (self._source_id,),
            )
            available = cursor.fetchone()[0]
            if len(ann_results) >= min(k, available):
                return ann_results

            cursor.execute("SET LOCAL enable_indexscan = off;")
            cursor.execute("SET LOCAL enable_bitmapscan = off;")
            cursor.execute(
                self._nearest_sql(),
                (query_vector, self._source_id, query_vector, k),
            )
            exact = cursor.fetchall()
            if len(exact) > len(ann_results):
                logging.info(
                    "Vector index under-returned for source %s (%d of %d); "
                    "used exact search instead.",
                    self._source_id, len(ann_results), min(k, available),
                )
                return exact
            return ann_results
        except Exception as e:
            # Never let the safety net take down the search it is protecting —
            # but roll back, or the aborted transaction poisons the connection
            # and every later search on this store returns nothing.
            logging.warning("Exact-search fallback failed: %s", e)
            try:
                cursor.connection.rollback()
            except Exception:
                # Connection already gone; nothing left to roll back.
                pass
            return ann_results
        finally:
            try:
                # RESET, not "= on": a deployment may disable these globally.
                cursor.execute("RESET enable_indexscan;")
                cursor.execute("RESET enable_bitmapscan;")
            except Exception:
                # Cursor/transaction already unusable; the settings die with it.
                pass

    def search_with_scores(
        self,
        question: str,
        k: int = 2,
        *args,
        score_threshold: float = None,
        **kwargs,
    ) -> List[tuple]:
        """Same search as :meth:`search`, pairing each hit with its similarity.

        The score is the cosine similarity (``1 - cosine_distance``) — the exact
        quantity ``score_threshold`` is compared against, so a caller can read a
        result's score and pick a threshold from it directly.
        """
        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 = self._nearest_sql()

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

            # An ANN index filters ``source_id`` *after* choosing candidates, so
            # a source holding a small share of the table can come back short —
            # or empty — no matter how the index is tuned. Raising probes /
            # ef_search only moves that threshold. When the result looks short,
            # redo the query exactly: correctness is worth one extra scan, and a
            # silent empty result reaches the model as "no documents exist".
            if len(results) < k:
                results = self._exact_search(cursor, query_vector, k, results)

            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 {}
                score = None if distance is None else 1.0 - float(distance)
                documents.append(
                    (Document(page_content=text, metadata=metadata), score)
                )

            return documents

        except Exception as e:
            logging.error(f"Error searching documents: {e}", exc_info=True)
            try:
                conn.rollback()
            except Exception:
                # Connection already gone; nothing left to roll back.
                pass
            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()