chore: import upstream snapshot with attribution

application/vectorstore/pgvector.py

@@ -0,0 +1,406 @@
+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()