[KM-436-437] edit knowledge handler pipeline

src/db/postgres/init_db.py

@@ -28,3 +28,38 @@ async def init_db():
         await conn.execute(text(
             "ALTER TABLE rooms ADD COLUMN IF NOT EXISTS status VARCHAR NOT NULL DEFAULT 'active'"
         ))
+
+        # HNSW index for fast approximate vector similarity search
+        # Only created when the embedding column has explicit dimensions (HNSW requirement).
+        # atttypmod > 0 means the vector column was created with a dimension (e.g. vector(1536));
+        # atttypmod = -1 means dimensionless — HNSW would fail with "column does not have dimensions".
+        await conn.execute(text("""
+            DO $$
+            BEGIN
+                IF EXISTS (
+                    SELECT FROM pg_attribute a
+                    JOIN pg_class c ON c.oid = a.attrelid
+                    WHERE c.relname = 'langchain_pg_embedding'
+                      AND a.attname = 'embedding'
+                      AND a.atttypmod > 0
+                ) THEN
+                    CREATE INDEX IF NOT EXISTS idx_langchain_pg_embedding_hnsw
+                    ON langchain_pg_embedding USING hnsw (embedding vector_cosine_ops);
+                END IF;
+            END $$
+        """))
+
+        # GIN index for FTS on schema chunks — only created if table exists
+        # (langchain_pg_embedding is created by PGVector on first use, not by create_all)
+        await conn.execute(text("""
+            DO $$
+            BEGIN
+                IF EXISTS (
+                    SELECT FROM information_schema.tables
+                    WHERE table_name = 'langchain_pg_embedding'
+                ) THEN
+                    CREATE INDEX IF NOT EXISTS idx_langchain_pg_embedding_fts
+                    ON langchain_pg_embedding USING GIN (to_tsvector('english', document));
+                END IF;
+            END $$
+        """))

src/document/document_service.py

@@ -1,8 +1,9 @@
 """Service for managing documents."""
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy import select, delete
+from sqlalchemy import select, delete, text
 from src.db.postgres.models import Document
+from src.db.postgres.connection import _pgvector_engine
 from src.storage.az_blob.az_blob import blob_storage
 from src.middlewares.logging import get_logger
 from typing import List, Optional
@@ -77,6 +78,21 @@ class DocumentService:
         # Delete from blob storage
         await blob_storage.delete_file(document.blob_name)
 
+        # Delete vector embeddings from pgvector (scoped to user + collection to avoid cross-user over-delete)
+        async with _pgvector_engine.begin() as conn:
+            await conn.execute(
+                text("""
+                    DELETE FROM langchain_pg_embedding
+                    WHERE cmetadata->>'user_id' = :user_id
+                      AND cmetadata->>'source_type' = 'document'
+                      AND cmetadata->'data'->>'document_id' = :doc_id
+                      AND collection_id = (
+                        SELECT uuid FROM langchain_pg_collection WHERE name = 'document_embeddings'
+                      )
+                """),
+                {"user_id": document.user_id, "doc_id": document_id},
+            )
+
         # Delete from database
         await db.execute(
             delete(Document).where(Document.id == document_id)

src/knowledge/parquet_service.py

@@ -0,0 +1,77 @@
+"""Parquet service — converts, uploads, downloads, and deletes Parquet files for CSV/XLSX.
+
+Parquet files are stored in Azure Blob alongside the original document using
+a deterministic naming convention based on document_id:
+
+  CSV:        {user_id}/{document_id}.parquet
+  XLSX sheet: {user_id}/{document_id}__{safe_sheet_name}.parquet
+
+This allows tabular.py to construct the correct blob name at retrieval time
+without needing to store it separately, and allows document_pipeline.py to
+delete all Parquet files for a document using a prefix delete.
+"""
+
+import io
+
+import pandas as pd
+
+from src.middlewares.logging import get_logger
+from src.storage.az_blob.az_blob import blob_storage
+
+logger = get_logger("parquet_service")
+
+
+def _safe_sheet_name(sheet_name: str) -> str:
+    return sheet_name.replace("/", "_").replace(" ", "_").replace("\\", "_")
+
+
+def parquet_blob_name(user_id: str, document_id: str, sheet_name: str | None = None) -> str:
+    """Construct deterministic Parquet blob name."""
+    if sheet_name:
+        return f"{user_id}/{document_id}__{_safe_sheet_name(sheet_name)}.parquet"
+    return f"{user_id}/{document_id}.parquet"
+
+
+def _to_parquet_bytes(df: pd.DataFrame) -> bytes:
+    buf = io.BytesIO()
+    df.to_parquet(buf, index=False)
+    return buf.getvalue()
+
+
+async def upload_parquet(
+    df: pd.DataFrame,
+    user_id: str,
+    document_id: str,
+    sheet_name: str | None = None,
+) -> str:
+    """Convert DataFrame to Parquet and upload to Azure Blob. Returns blob_name."""
+    blob_name = parquet_blob_name(user_id, document_id, sheet_name)
+    parquet_bytes = _to_parquet_bytes(df)
+    await blob_storage.upload_bytes(parquet_bytes, blob_name)
+    logger.info(f"Uploaded Parquet {blob_name} ({len(parquet_bytes)} bytes)")
+    return blob_name
+
+
+async def download_parquet(
+    user_id: str,
+    document_id: str,
+    sheet_name: str | None = None,
+) -> pd.DataFrame:
+    """Download Parquet from Azure Blob and return as DataFrame."""
+    blob_name = parquet_blob_name(user_id, document_id, sheet_name)
+    content = await blob_storage.download_file(blob_name)
+    df = pd.read_parquet(io.BytesIO(content))
+    logger.info(f"Downloaded Parquet {blob_name}: {len(df)} rows, {len(df.columns)} columns")
+    return df
+
+
+async def delete_document_parquets(user_id: str, document_id: str) -> int:
+    """Delete all Parquet files for a document (CSV = 1 file, XLSX = one per sheet).
+
+    Uses prefix delete: {user_id}/{document_id} matches all Parquet variants
+    for this document without touching the original blob (which uses a random UUID name).
+    """
+    prefix = f"{user_id}/{document_id}"
+    deleted = await blob_storage.delete_blobs_with_prefix(prefix)
+    logger.info(f"Deleted {deleted} Parquet file(s) for document {document_id}")
+    return deleted

src/knowledge/processing_service.py

@@ -7,7 +7,9 @@ from src.storage.az_blob.az_blob import blob_storage
 from src.db.postgres.models import Document as DBDocument
 from sqlalchemy.ext.asyncio import AsyncSession
 from src.middlewares.logging import get_logger
+from src.knowledge.parquet_service import upload_parquet
 from typing import List
+from datetime import datetime, timezone, timedelta
 import sys
 import docx
 import pandas as pd
@@ -15,6 +17,8 @@ import pytesseract
 from pdf2image import convert_from_bytes
 from io import BytesIO
 
+_JAKARTA_TZ = timezone(timedelta(hours=7))
+
 logger = get_logger("knowledge_processing")
 
 
@@ -41,9 +45,9 @@ class KnowledgeProcessingService:
             if db_doc.file_type == "pdf":
                 documents = await self._build_pdf_documents(content, db_doc)
             elif db_doc.file_type == "csv":
-                documents = self._build_csv_documents(content, db_doc)
+                documents = await self._build_csv_documents(content, db_doc)
             elif db_doc.file_type == "xlsx":
-                documents = self._build_excel_documents(content, db_doc)
+                documents = await self._build_excel_documents(content, db_doc)
             else:
                 text = self._extract_text(content, db_doc.file_type)
                 if not text.strip():
@@ -55,6 +59,7 @@ class KnowledgeProcessingService:
                         metadata={
                             "user_id": db_doc.user_id,
                             "source_type": "document",
+                            "updated_at": datetime.now(_JAKARTA_TZ).isoformat(),
                             "data": {
                                 "document_id": db_doc.id,
                                 "filename": db_doc.filename,
@@ -103,6 +108,7 @@ class KnowledgeProcessingService:
                     metadata={
                         "user_id": db_doc.user_id,
                         "source_type": "document",
+                        "updated_at": datetime.now(_JAKARTA_TZ).isoformat(),
                         "data": {
                             "document_id": db_doc.id,
                             "filename": db_doc.filename,
@@ -150,6 +156,7 @@ class KnowledgeProcessingService:
                 metadata={
                     "user_id": db_doc.user_id,
                     "source_type": "document",
+                    "updated_at": datetime.now(_JAKARTA_TZ).isoformat(),
                     "data": {
                         "document_id": db_doc.id,
                         "filename": db_doc.filename,
@@ -162,18 +169,25 @@ class KnowledgeProcessingService:
             ))
         return documents
 
-    def _build_csv_documents(self, content: bytes, db_doc: DBDocument) -> List[LangChainDocument]:
-        """Profile each column of a CSV file."""
+    async def _build_csv_documents(self, content: bytes, db_doc: DBDocument) -> List[LangChainDocument]:
+        """Profile each column of a CSV file and upload Parquet to Azure Blob."""
         df = pd.read_csv(BytesIO(content))
+        await upload_parquet(df, db_doc.user_id, db_doc.id)
+        logger.info(f"Uploaded Parquet for CSV {db_doc.id}")
         return self._profile_dataframe(df, db_doc.filename, db_doc)
 
-    def _build_excel_documents(self, content: bytes, db_doc: DBDocument) -> List[LangChainDocument]:
-        """Profile each column of every sheet in an Excel file."""
+    async def _build_excel_documents(self, content: bytes, db_doc: DBDocument) -> List[LangChainDocument]:
+        """Profile each column of every sheet in an Excel file and upload one Parquet per sheet."""
         sheets = pd.read_excel(BytesIO(content), sheet_name=None)
         documents = []
         for sheet_name, df in sheets.items():
             source_name = f"{db_doc.filename} / sheet: {sheet_name}"
-            documents.extend(self._profile_dataframe(df, source_name, db_doc))
+            docs = self._profile_dataframe(df, source_name, db_doc)
+            for doc in docs:
+                doc.metadata["data"]["sheet_name"] = sheet_name
+            documents.extend(docs)
+            await upload_parquet(df, db_doc.user_id, db_doc.id, sheet_name)
+            logger.info(f"Uploaded Parquet for sheet '{sheet_name}' of {db_doc.id}")
         return documents
 
     def _extract_text(self, content: bytes, file_type: str) -> str:

src/models/sql_query.py

@@ -0,0 +1,8 @@
+"""Structured output model for LLM-generated SQL queries."""
+
+from pydantic import BaseModel, Field
+
+
+class SQLQuery(BaseModel):
+    sql: str = Field(description="A single SQL SELECT statement. No markdown, no explanation inline.")
+    reasoning: str = Field(description="One sentence: what this query answers.")

src/models/structured_output.py

@@ -19,3 +19,7 @@ class IntentClassification(BaseModel):
         default="",
         description="Direct response if no search needed (for greetings, etc.)"
     )
+    source_hint: str = Field(
+        default="both",
+        description="Which sources to search: 'document' (PDF/DOCX/TXT), 'schema' (DB/CSV/XLSX), or 'both'"
+    )

src/pipeline/db_pipeline/db_pipeline_service.py

@@ -148,7 +148,7 @@ class DbPipelineService:
             engine.dispose()
 
     def _to_document(
-        self, user_id: str, table_name: str, entry: dict, updated_at: str
+        self, user_id: str, client_id: str, table_name: str, entry: dict, updated_at: str
     ) -> LangChainDocument:
         col = entry["col"]
         return LangChainDocument(
@@ -156,6 +156,7 @@ class DbPipelineService:
             metadata={
                 "user_id": user_id,
                 "source_type": "database",
+                "database_client_id": client_id,
                 "updated_at": updated_at,
                 "data": {
                     "table_name": table_name,
@@ -170,6 +171,7 @@ class DbPipelineService:
     async def run(
         self,
         user_id: str,
+        client_id: str,
         engine: Engine,
         exclude_tables: Optional[frozenset[str]] = None,
     ) -> int:
@@ -181,35 +183,50 @@ class DbPipelineService:
         vector_store = get_vector_store()
         logger.info("db pipeline start", user_id=user_id)
 
+        # Profile first — if this fails, old embeddings are untouched
+        schema = await asyncio.to_thread(get_schema, engine, exclude_tables)
+
+        updated_at = datetime.now(timezone(timedelta(hours=7))).isoformat()
+        all_docs: list = []
+        for table_name, columns in schema.items():
+            logger.info("profiling table", table=table_name, columns=len(columns))
+            entries = await asyncio.to_thread(profile_table, engine, table_name, columns)
+            docs = [self._to_document(user_id, client_id, table_name, e, updated_at) for e in entries]
+            all_docs.extend(docs)
+            logger.info("profiled table", table=table_name, count=len(docs))
+
+        # Insert new chunks first; only delete stale chunks after the insert succeeds.
+        # Prevents data loss if aadd_documents fails — old embeddings stay queryable
+        # until they're proven replaceable. Stale rows are identified by an older
+        # updated_at than this run.
+        if not all_docs:
+            logger.warning(
+                "no docs produced from schema; skipping delete to preserve existing embeddings",
+                user_id=user_id,
+                client_id=client_id,
+            )
+            return 0
+
+        await vector_store.aadd_documents(all_docs)
+
         async with _pgvector_engine.begin() as conn:
             result = await conn.execute(
                 text(
                     "DELETE FROM langchain_pg_embedding "
                     "WHERE cmetadata->>'user_id' = :user_id "
                     "  AND cmetadata->>'source_type' = 'database' "
+                    "  AND cmetadata->>'database_client_id' = :client_id "
+                    "  AND cmetadata->>'updated_at' < :updated_at "
                     "  AND collection_id = ("
                     "    SELECT uuid FROM langchain_pg_collection WHERE name = 'document_embeddings'"
                     "  )"
                 ),
-                {"user_id": user_id},
+                {"user_id": user_id, "client_id": client_id, "updated_at": updated_at},
             )
-            logger.info("cleared old db embeddings", user_id=user_id, deleted=result.rowcount)
-
-        schema = await asyncio.to_thread(get_schema, engine, exclude_tables)
+            logger.info("cleared stale db embeddings", user_id=user_id, deleted=result.rowcount)
 
-        updated_at = datetime.now(timezone(timedelta(hours=7))).isoformat()
-        total = 0
-        for table_name, columns in schema.items():
-            logger.info("profiling table", table=table_name, columns=len(columns))
-            entries = await asyncio.to_thread(profile_table, engine, table_name, columns)
-            docs = [self._to_document(user_id, table_name, e, updated_at) for e in entries]
-            if docs:
-                await vector_store.aadd_documents(docs)
-                total += len(docs)
-                logger.info("ingested chunks", table=table_name, count=len(docs))
-
-        logger.info("db pipeline complete", user_id=user_id, total=total)
-        return total
+        logger.info("db pipeline complete", user_id=user_id, total=len(all_docs))
+        return len(all_docs)
 
 
 db_pipeline_service = DbPipelineService()

src/pipeline/document_pipeline/document_pipeline.py

@@ -5,6 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from src.document.document_service import document_service
 from src.knowledge.processing_service import knowledge_processor
+from src.knowledge.parquet_service import delete_document_parquets
 from src.middlewares.logging import get_logger
 from src.storage.az_blob.az_blob import blob_storage
 
@@ -32,7 +33,7 @@ class DocumentPipeline:
         if file_type not in SUPPORTED_FILE_TYPES:
             raise HTTPException(
                 status_code=400,
-                detail=f"Unsupported file type. Supported: {SUPPORTED_FILE_TYPES}",
+                detail=f"Unsupported file type. Supported: {', '.join(SUPPORTED_FILE_TYPES)}",
             )
 
         blob_name = await blob_storage.upload_file(content, file.filename, user_id)
@@ -81,6 +82,9 @@ class DocumentPipeline:
 
         await document_service.delete_document(db, document_id)
 
+        if document.file_type in ("csv", "xlsx"):
+            await delete_document_parquets(user_id, document_id)
+
         logger.info(f"Deleted document {document_id} for user {user_id}")
         return {"document_id": document_id}
 

src/storage/az_blob/az_blob.py

@@ -57,6 +57,22 @@ class AzureBlobStorage:
             logger.error(f"Failed to download blob {blob_name}", error=str(e))
             raise
 
+    async def upload_bytes(self, content: bytes, blob_name: str) -> str:
+        """Upload bytes to Azure Blob Storage using a specific blob name.
+
+        Unlike upload_file(), this does not generate a UUID name — caller controls the blob_name.
+        Used for Parquet files where the name must be deterministic (derived from document_id).
+        """
+        try:
+            async with self._get_blob_client(blob_name) as blob_client:
+                logger.info(f"Uploading bytes to blob {blob_name}")
+                await blob_client.upload_blob(content, overwrite=True)
+            logger.info(f"Successfully uploaded {blob_name}")
+            return blob_name
+        except Exception as e:
+            logger.error(f"Failed to upload bytes to {blob_name}", error=str(e))
+            raise
+
     async def delete_file(self, blob_name: str) -> bool:
         """Delete file from Azure Blob Storage."""
         try:
@@ -71,6 +87,24 @@ class AzureBlobStorage:
             logger.error(f"Failed to delete blob {blob_name}", error=str(e))
             return False
 
+    async def delete_blobs_with_prefix(self, prefix: str) -> int:
+        """Delete all blobs whose name starts with prefix. Returns count deleted.
+
+        Used to delete all Parquet files for a document in one call.
+        """
+        from azure.storage.blob.aio import ContainerClient
+        container_url = f"{self.account_url}/{self.container_name}?{self.sas_token}"
+        deleted = 0
+        try:
+            async with ContainerClient.from_container_url(container_url) as container:
+                async for blob in container.list_blobs(name_starts_with=prefix):
+                    await container.delete_blob(blob.name)
+                    deleted += 1
+            logger.info(f"Deleted {deleted} blobs with prefix {prefix}")
+        except Exception as e:
+            logger.error(f"Failed to delete blobs with prefix {prefix}", error=str(e))
+        return deleted
+
 
 # Singleton instance
 blob_storage = AzureBlobStorage()