✨ feat: Optimize PostgreSQL Queries and Add Optional Detailed Query Logging

- Introduced a new environment variable `DEBUG_PGVECTOR_QUERIES` to enable detailed logging of pgvector operations.
- Implemented query logging setup in the ExtendedPgVector class, capturing execution time and parameters for relevant queries.
- Updated README.md to document the new environment variable.

Author: danny-avila

README.md

@@ -61,6 +61,7 @@ The following environment variables are required to run the application:
 - `RAG_UPLOAD_DIR`: (Optional) The directory where uploaded files are stored. Default value is "./uploads/".
 - `PDF_EXTRACT_IMAGES`: (Optional) A boolean value indicating whether to extract images from PDF files. Default value is "False".
 - `DEBUG_RAG_API`: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes
+- `DEBUG_PGVECTOR_QUERIES`: (Optional) Set to "True" to enable detailed PostgreSQL query logging for pgvector operations. Useful for debugging performance issues with vector database queries.
 - `CONSOLE_JSON`: (Optional) Set to "True" to log as json for Cloud Logging aggregations
 - `EMBEDDINGS_PROVIDER`: (Optional) either "openai", "bedrock", "azure", "huggingface", "huggingfacetei", "vertexai", or "ollama", where "huggingface" uses sentence_transformers; defaults to "openai"
 - `EMBEDDINGS_MODEL`: (Optional) Set a valid embeddings model to use from the configured provider.

app/services/vector_store/extended_pg_vector.py

@@ -1,19 +1,73 @@
+import os
+import time
+import logging
 from typing import Optional
-
+from sqlalchemy import event
 from sqlalchemy import delete
 from sqlalchemy.orm import Session
+from sqlalchemy.engine import Engine
 from langchain_core.documents import Document
 from langchain_community.vectorstores.pgvector import PGVector
 
+
 class ExtendedPgVector(PGVector):
+    _query_logging_setup = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.setup_query_logging()
+
+    def setup_query_logging(self):
+        """Enable query logging for this vector store only if DEBUG_PGVECTOR_QUERIES is set"""
+        # Only setup logging if the environment variable is set to a truthy value
+        debug_queries = os.getenv("DEBUG_PGVECTOR_QUERIES", "").lower()
+        if debug_queries not in ["true", "1", "yes", "on"]:
+            return
+
+        # Only setup once per class
+        if ExtendedPgVector._query_logging_setup:
+            return
+
+        logger = logging.getLogger("pgvector.queries")
+        logger.setLevel(logging.INFO)
+
+        # Create handler if it doesn't exist
+        if not logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter("%(asctime)s - PGVECTOR QUERY - %(message)s")
+            handler.setFormatter(formatter)
+            logger.addHandler(handler)
+
+        @event.listens_for(Engine, "before_cursor_execute")
+        def receive_before_cursor_execute(
+            conn, cursor, statement, parameters, context, executemany
+        ):
+            if "langchain_pg_embedding" in statement:
+                context._query_start_time = time.time()
+                logger.info(f"STARTING QUERY: {statement}")
+                logger.info(f"PARAMETERS: {parameters}")
+
+        @event.listens_for(Engine, "after_cursor_execute")
+        def receive_after_cursor_execute(
+            conn, cursor, statement, parameters, context, executemany
+        ):
+            if "langchain_pg_embedding" in statement:
+                total = time.time() - context._query_start_time
+                logger.info(f"COMPLETED QUERY in {total:.4f}s")
+                logger.info("-" * 50)
+
+        ExtendedPgVector._query_logging_setup = True
+
     def get_all_ids(self) -> list[str]:
         with Session(self._bind) as session:
             results = session.query(self.EmbeddingStore.custom_id).all()
             return [result[0] for result in results if result[0] is not None]
-        
+
     def get_filtered_ids(self, ids: list[str]) -> list[str]:
         with Session(self._bind) as session:
-            query = session.query(self.EmbeddingStore.custom_id).filter(self.EmbeddingStore.custom_id.in_(ids))
+            query = session.query(self.EmbeddingStore.custom_id).filter(
+                self.EmbeddingStore.custom_id.in_(ids)
+            )
             results = query.all()
             return [result[0] for result in results if result[0] is not None]
 
@@ -45,7 +99,9 @@ def _delete_multiple(
                     if not collection:
                         self.logger.warning("Collection not found")
                         return
-                    stmt = stmt.where(self.EmbeddingStore.collection_id == collection.uuid)
+                    stmt = stmt.where(
+                        self.EmbeddingStore.collection_id == collection.uuid
+                    )
                 stmt = stmt.where(self.EmbeddingStore.custom_id.in_(ids))
                 session.execute(stmt)
-            session.commit()
+            session.commit()