Merge dev_new to main (improve retriever and add querying) (#13)

- [noticket] add gitignore (c87f27f458854787102eebc110b178c2572ef8a6)
- [NOTICKET]: add document pipeline, simplify document API (fb871f3d7db9c5a9c95a012761c4180a432e7e4b)
- [NOTICKET]: update folder document_pipelines after pipelines (a4cf97ab3890e35831c0043a1d411c17d84f374b)
- [NOTICKET][DB] refactor code to new repo (7f3bb9782bb9b6ace69e6be35377c12757a2f81d)
- [KM-441] add mean and median (9b5933420821187e73bb72a896e28d7adf6d6137)
- [NOTICKET] new metadata format for cleaner code (6b590d94c398ac4966a0e376f2809aa52c5ba022)
- update document (5a69e0eecf80dd807565dbc369374e41450c37fc)
- delete duplicate file (3848d7b262745b758069e3ab8089437ced42a45f)
- edit document for new pipeline (425e021053c4de65e2598b0f5554b7c1044a904c)
- [NOTICKET]: add CSV and XLSX file type (31920c3bc6f9faf9e07095a02fc1faf08adc246e)
- [DB] fix/rename db_pipeline.py (d913315c356d741cc7a5a5b1f8e36ba41397f2ab)
- [NOTICKET][DB] menyesuaikan format struktur db_pipeline sesuai dengan file lain (e13a9017e0f7fd26ef9b8ad5981d6c0799c01661)
- [NOTICKET][DB] pisahin db credential ke folder model. add ingestion endpoint at db_client to use db pipeline. add router db_client di main. (347a73aa5597cb28d1f43940eb60db619e4a10be)
- [NOTICKET]: use tesseract for extract PDF (6b9a13d417f6ffffed61a44567a3acbee822e173)
- [NOTICKET]: add Tesseract and Poppler binaries via Git LFS (0a9101a16c713c3facd8c06279da1f01403306a0)
- [NOTICKET]: update uv.lock (bb79f64b69c3bdcb2c98867c5d4945bf2311394d)
- [NOTICKET][DB] update credential & databaseclient. update settings (0e079550f8524da66df365638167dfe42b70ae62)
- [NOTICKET] update settings (65a5c6b1c1539de5d66e700aa5bef011e3c2b2f6)
- [KM-437][DB] add mysql, sqlserver, bigquery, snowflake connections (43539293eab9328d1d6f22f2f837c0a0085485a5)
- [NOTICKET]: adjusted pyproject.toml for OCR PDF (a00e2ad5306306ff73c30d085f126fab7bd1e78c)
- [NOTICKET]: fix merge conflict (6c8734607da68151ab574d0d186b17ac0d8ec460)
- [NOTICKET][DB] fix mysql pipeline (060c8cc81d4de2e85d2f6b4d3de178307e211ad0)
- [NOTICKET] edit imports (b145c06e354030dc38a550cd53030ba4864870ae)
- [NOTICKET] minor code refactor (52415b6a52f08900d27b548459a15c1113c1b892)
- [NOTICKET] add duplicate check for storing database (d310770ffcbcf75cb6b952f86bc6947ab2c7bb6a)
- [NOTICKET][DB] add supported dbtype for frontend (a531fcc7f53fe56a634ebba3b0fd4ff7ad51a33d)
- [NOTICKET]: add doctypes endpoint & 10MB file size limit (9debae56609158434faf47bf2a4a7b70c4137290)
- [NOTICKET]: add comments to flag that file type lists must stay in sync (023b7cfeb1dd6f6c54df748feb1130e4a994bc26)
- [NOTICKET]: add to gitignore (bbc8c584ef676c15d913d04933b9b8ff2b8d98d8)
- Merge branch 'main' of https://huggingface.co/spaces/DataEyond/Agentic-Service-Data-Eyond into dev_new (7757da18fa23f89a4d731cf2648727a97c49b88c)
- Merge branch 'main' of https://huggingface.co/spaces/DataEyond/Agentic-Service-Data-Eyond into dev_new (9c090a04cf8a78322836d04e245632bbf1f04555)
- Merge branch 'main' of https://huggingface.co/spaces/DataEyond/Agentic-Service-Data-Eyond into dev_new (5398fec4138f7f2b944c0a3ddbdf1f858bac7e16)
- Merge branch 'dev_new' of https://huggingface.co/spaces/DataEyond/Agentic-Service-Data-Eyond into dev_new (20bf3f8fa4393217cb50f9e330196b26bb138fa1)
- [NOTICKET] add total token logging (b9703fc5cfad1ec51ef67497559ef4a5235ea3d5)
- [NOTICKET] add updated_at field for metadata & delete old embedding before appending (cb5ab3273516de824aa254ff7c9fed0c5f8098aa)
- [NO TICKET][document]: add updated_at on metadata (d2f7a4833e1069abc3d3831381661def3af4ad8a)
- [NO TICKET][document]: delete vector embedding on table langchain_pg_embedding if user delete document on knowledge (ac3d8c19b4edb19fad7630b7ec5dabc39f9948ac)
- [NOTICKET][document]: make a clean output to status error unsupported file type (2814813f836a095af5b20552ec047cbc29c95821)
- [KM-438][KM-439] framework for knowledge retriever (d1e12641c43e217c684389dd6697f29ffc489226)
- Merge branch 'main' of https://huggingface.co/spaces/DataEyond/Agentic-Service-Data-Eyond into dev_new (a701ac3732bf0eb3da2b13f1801c89537986174e)
- [NOTICKET] fix single source to multiple sources (589ca324a5fec304a2d3652b65311e494c5a44b9)
- [NOTICKET] fixed multiple sources (e9f2a26317988ef59031c32726cae793ad135755)
- [KM-507] add multiple retrieval method to compare (dense, mmr, bm25, hybrid) (ac6b78d1a5448aab8a210416374bd3ffad4300dd)
- [KM-507] add changes to methods (82186504125f65d3767fee5371451847f435a9ea)
- [NOTICKET] add db_client for querying (e49db6011cc3be00efeb0807c02db0bbffa56999)
- add to gitignore (83ed7447e869f2d884d9cd341b7e20bac00f48cd)
- [KM-507] add different methods, now using dense cosine (145bca392470b81e16060feb040ab652f7185976)
- [KM-512] create folder for querying from bd/tabular docs (2c8a3e89bb68fe046beac5dc9934a8d6daaba15d)
- [NOTICKET] minor fix in chat.py, add package for query, change schema used to hybrid (cosine+bm25) (15cd3a7f887abe8994b45d80d4d7ade688120b88)
- [KM-512] add Pydantic model the LLM fills via function calling in sql_query, and add same signature for db and tabular (220f59eb60b88891e476c989b9549cc180cc08b7)
- [NOTICKET] rename file name, updated after uv sync (948d6ddaa2b4bdafef45a219169832d6c7b602bc)
- [NOTICKET] update .gitignore (240251c4050c9a6457522d0024e87f99dbfbbb13)
- Merge branch 'dev_new' of https://huggingface.co/spaces/DataEyond/Agentic-Service-Data-Eyond into dev_new (29efec6750e81688c99f62f52b0fdc137da8fee8)
- [KM-513][document] add convert to parquet if type file is XLSX and CSV (770f26b143a795987628fbdb4716ee8c574a6e5c)
- add to gitignore (1fef470b29037577cde6246ec7f417f164fed72e)
- [KM-512] connect query executor to user question. add logging for db_executor (abc494f97718a67af0d237ebac88fb240f3b46d3)
- [NOTICKET] fix delete, now can filter by user (f273db057a8158fe7eac526688ffd52d9ebaf911)
- [NOTICKET] db_executor: CTE DML check now walks entire AST root, schema: cast instead of string interpolation (bd2b1d9ddf6664c2739563fbd04eaca279771ec0)
- [NOTICKET] fix-revert string change (110ee3438ef171e177c53f38b437e6dc8e927a42)
- [KM-520] Integrate db query executor pipeline with existing rag retrieve pipeline (a25febe2e68afe14a0c75352fe746dc9d0157643)
- [KM-516][KM-517] add new feature; ai can now see table & column names that have fk relationship with retrieved result (f86da27bdc2a60f0b38b61db90008094ee4d7146)
- [NOTICKET] fix query now use orchestrator msg, rework db pipeline replace ingestion logic (be9bbd9d4451e320e923627eca36f1442cae1932)
- [NOTICKET] edit settings (9e16c227411e7093aea1bd09b79c5bf60fd0431b)
- [KM 436-439] adjust endpoint for new features (0707f2ba972cd3b62aadd4cee8b445ba9aeb6ede)
- [KM-436-437] edit knowledge handler pipeline (767625e5fb1b1371135439c7b9c6cf1e6389d218)
- [KM-438-439] add retriever feature (ba550a5a0e6cd7a64dc471711116e4ce81402049)


Co-authored-by: Sofhia Az-Zahra <sofhiaazzhr@users.noreply.huggingface.co>

.gitignore

@@ -28,6 +28,8 @@ test/users/user_accounts.csv
 
 CLAUDE.md
 
+/experiments
+src/rag/experiments/
 erd/
 playground/
 playground_retriever.py

pyproject.toml

@@ -79,6 +79,8 @@ dependencies = [
     "jsonpatch>=1.33",
     "pymongo>=4.14.0",
     "psycopg2>=2.9.11",
+    # --- SQL parsing / guardrails ---
+    "sqlglot>=25.0.0",
     # --- User-DB connectors (db_pipeline) ---
     "pymysql>=1.1.1",
     "pymssql>=2.3.0",

src/agents/orchestration.py

@@ -35,6 +35,11 @@ Intent Routing:
 - greeting -> needs_search=False, direct_response="Hello! How can I assist you today?"
 - goodbye -> needs_search=False, direct_response="Goodbye! Have a great day!"
 - other -> needs_search=True, search_query=<standalone rewritten query>
+
+Source Routing (set source_hint):
+- Columns, tables, sheets, data types, schema, row counts, statistics -> source_hint=schema
+- Document content, paragraphs, reports, articles, text -> source_hint=document
+- Unclear or spans both -> source_hint=both
 """),
             MessagesPlaceholder(variable_name="history"),
             ("user", "{message}")

src/api/v1/chat.py

@@ -9,6 +9,9 @@ from src.db.postgres.models import ChatMessage, MessageSource
 from src.agents.orchestration import orchestrator
 from src.agents.chatbot import chatbot
 from src.rag.retriever import retriever
+from src.rag.base import RetrievalResult
+from src.query.query_executor import query_executor
+from src.query.base import QueryResult
 from src.db.redis.connection import get_redis
 from src.config.settings import settings
 from src.middlewares.logging import get_logger, log_execution
@@ -61,7 +64,7 @@ def _extract_sources(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     seen = set()
     sources = []
     for result in results:
-        if "document_id" in result["metadata"]["data"]:
+        if "document_id" in result["metadata"].get("data", {}):
             meta = result["metadata"]
             key = (meta.get("data", {}).get("document_id"), meta.get("data", {}).get("page_label"))
             if key not in seen:
@@ -88,6 +91,22 @@ def _extract_sources(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     return sources
 
 
+def _format_query_results(results: list[QueryResult]) -> str:
+    if not results:
+        return ""
+    lines = []
+    for r in results:
+        name = r.metadata.get("client_name", r.source_id)
+        lines.append(f"[Query result — {name}, tables: {r.table_or_file}]")
+        lines.append(f"SQL: {r.metadata.get('sql', '')}")
+        if r.columns and r.rows:
+            lines.append(" | ".join(r.columns))
+            for row in r.rows[:20]:
+                lines.append(" | ".join(str(row.get(c, "")) for c in r.columns))
+        lines.append(f"({r.row_count} rows total)\n")
+    return "\n".join(lines)
+
+
 async def get_cached_response(redis, cache_key: str) -> Optional[str]:
     cached = await redis.get(cache_key)
     if cached:
@@ -182,16 +201,25 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
 
             if not intent_result.get("needs_search"):
                 retrieval_task.cancel()
+                try:
+                    await retrieval_task
+                except asyncio.CancelledError:
+                    pass
                 raw_results = []
             else:
                 search_query = intent_result.get("search_query", request.message)
                 logger.info(f"Searching for: {search_query}")
                 if search_query != request.message:
                     retrieval_task.cancel()
+                    try:
+                        await retrieval_task
+                    except asyncio.CancelledError:
+                        pass
                     raw_results = await retriever.retrieve(
                         query=search_query,
                         user_id=request.user_id,
                         db=db,
+                        source_hint=intent_result.get("source_hint", "both"),
                     )
                 else:
                     raw_results = await retrieval_task
@@ -199,6 +227,27 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
             context = _format_context(raw_results)
             sources = _extract_sources(raw_results)
 
+            source_hint = intent_result.get("source_hint", "both")
+            if source_hint in ("schema", "both"):
+                retrieval_objects = [
+                    RetrievalResult(
+                        content=r["content"],
+                        metadata=r["metadata"],
+                        score=0.0,
+                        source_type=r["metadata"].get("source_type", ""),
+                    )
+                    for r in raw_results
+                ]
+                query_results = await query_executor.execute(
+                    results=retrieval_objects,
+                    user_id=request.user_id,
+                    db=db,
+                    question=intent_result.get("search_query") or request.message,
+                )
+                query_context = _format_query_results(query_results)
+                if query_context:
+                    context = query_context + "\n\n" + context
+
         # Step 3: Direct response for greetings / non-document intents
         if intent_result.get("direct_response"):
             response = intent_result["direct_response"]

src/api/v1/db_client.py

@@ -458,7 +458,7 @@ async def ingest_database_client(
             db_type=client.db_type,
             credentials=creds,
         ) as engine:
-            total = await db_pipeline_service.run(user_id=user_id, engine=engine)
+            total = await db_pipeline_service.run(user_id=user_id, client_id=client_id, engine=engine)
     except NotImplementedError as e:
         raise HTTPException(status_code=status.HTTP_501_NOT_IMPLEMENTED, detail=str(e))
     except Exception as e:

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/query/base.py

@@ -0,0 +1,32 @@
+"""Shared contract for query executors."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src.rag.base import RetrievalResult
+
+
+@dataclass
+class QueryResult:
+    source_type: str        # "database" or "document"
+    source_id: str          # database_client_id or document_id
+    table_or_file: str
+    columns: list[str]
+    rows: list[dict]
+    row_count: int
+    metadata: dict = field(default_factory=dict)
+    # metadata should include "column_types": {"col_name": "dtype"} when available
+
+
+class BaseExecutor(ABC):
+    @abstractmethod
+    async def execute(
+        self,
+        results: list[RetrievalResult],
+        user_id: str,
+        db: AsyncSession,
+        question: str,
+        limit: int = 100,
+    ) -> list[QueryResult]: ...

src/query/executors/db_executor.py

@@ -0,0 +1,409 @@
+"""Executor for registered database sources (source_type="database").
+
+Flow per (client_id, question):
+  1. Collect all relevant (table_name, column_name) pairs from retrieval results.
+  2. Fetch the FULL schema for those tables from PGVector (not just top-k columns).
+  3. Build a schema context string and send to LLM → structured SQLQuery output.
+  4. Validate via sqlglot: SELECT-only, schema-grounded, LIMIT enforced.
+  5. Execute on the user's DB via engine_scope + asyncio.to_thread.
+  6. Return QueryResult per client_id (may span multiple tables via JOINs).
+
+Supported db_types: postgres, supabase, mysql.
+Other types are skipped with a warning — they do not raise.
+"""
+
+import asyncio
+from collections import defaultdict
+from typing import Any
+
+import sqlglot
+import sqlglot.expressions as exp
+import tiktoken
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import AzureChatOpenAI
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src.config.settings import settings
+from src.database_client.database_client_service import database_client_service
+from src.db.postgres.connection import _pgvector_engine
+from src.middlewares.logging import get_logger
+from src.models.sql_query import SQLQuery
+from src.pipeline.db_pipeline import db_pipeline_service
+from src.query.base import BaseExecutor, QueryResult
+from src.rag.base import RetrievalResult
+from src.utils.db_credential_encryption import decrypt_credentials_dict
+
+logger = get_logger("db_executor")
+
+_enc = tiktoken.get_encoding("cl100k_base")
+
+_SUPPORTED_DB_TYPES = {"postgres", "supabase", "mysql"}
+_MAX_RETRIES = 3
+_MAX_LIMIT = 500
+
+_SQL_SYSTEM_PROMPT = """\
+You are a SQL data analyst working with a user's database.
+Generate a single SQL SELECT statement that answers the user's question.
+
+Database dialect: {dialect}
+
+Rules:
+- ONLY reference tables and columns listed in the schema below. Do not invent names.
+- Always include a LIMIT clause (max {limit}).
+- Do not use DELETE, UPDATE, INSERT, DROP, TRUNCATE, ALTER, CREATE, or any DDL.
+- Prefer explicit JOINs over subqueries when combining tables.
+- For aggregations, always alias the result column (e.g. COUNT(*) AS order_count).
+- For date filtering, use dialect-appropriate functions ({dialect} syntax).
+
+Schema:
+{schema}
+
+{error_section}"""
+
+
+class DbExecutor(BaseExecutor):
+    def __init__(self) -> None:
+        self._llm = AzureChatOpenAI(
+            azure_deployment=settings.azureai_deployment_name_4o,
+            openai_api_version=settings.azureai_api_version_4o,
+            azure_endpoint=settings.azureai_endpoint_url_4o,
+            api_key=settings.azureai_api_key_4o,
+            temperature=0,
+        )
+        self._prompt = ChatPromptTemplate.from_messages([
+            ("system", _SQL_SYSTEM_PROMPT),
+            ("human", "{question}"),
+        ])
+        self._chain = self._prompt | self._llm.with_structured_output(SQLQuery)
+
+    # ------------------------------------------------------------------
+    # Public interface
+    # ------------------------------------------------------------------
+
+    async def execute(
+        self,
+        results: list[RetrievalResult],
+        user_id: str,
+        db: AsyncSession,
+        question: str,
+        limit: int = 100,
+    ) -> list[QueryResult]:
+        db_results = [r for r in results if r.source_type == "database"]
+        if not db_results:
+            return []
+
+        # Group by client_id — one SQL generation + execution pass per client
+        by_client: dict[str, list[RetrievalResult]] = defaultdict(list)
+        for r in db_results:
+            client_id = r.metadata.get("database_client_id", "")
+            if client_id:
+                by_client[client_id].append(r)
+            else:
+                logger.warning("db result missing database_client_id, skipping")
+
+        query_results: list[QueryResult] = []
+        for client_id, client_results in by_client.items():
+            try:
+                qr = await self._execute_for_client(client_id, client_results, user_id, db, question, limit)
+                if qr:
+                    query_results.append(qr)
+            except Exception as e:
+                logger.error("db executor failed for client", client_id=client_id, error=str(e))
+
+        return query_results
+
+    # ------------------------------------------------------------------
+    # Per-client execution
+    # ------------------------------------------------------------------
+
+    async def _execute_for_client(
+        self,
+        client_id: str,
+        results: list[RetrievalResult],
+        user_id: str,
+        db: AsyncSession,
+        question: str,
+        limit: int,
+    ) -> QueryResult | None:
+        client = await database_client_service.get(db, client_id)
+        if not client:
+            logger.warning("database client not found", client_id=client_id)
+            return None
+        if client.user_id != user_id:
+            logger.warning("client ownership mismatch", client_id=client_id)
+            return None
+        if client.db_type not in _SUPPORTED_DB_TYPES:
+            logger.warning("unsupported db_type for query execution", db_type=client.db_type)
+            return None
+
+        # Distinct table names from retrieval results, expanded via FK relationships
+        table_names = list({
+            r.metadata.get("data", {}).get("table_name")
+            for r in results
+            if r.metadata.get("data", {}).get("table_name")
+        })
+        table_names = await self._expand_with_fk_tables(client_id, user_id, table_names)
+
+        full_schema = await self._fetch_full_schema(client_id, table_names, user_id)
+        if not full_schema:
+            logger.warning("no schema found in vector store", client_id=client_id, tables=table_names)
+            return None
+
+        schema_ctx = self._build_schema_context(full_schema)
+        capped_limit = min(limit, _MAX_LIMIT)
+        dialect = client.db_type
+
+        # SQL generation with retry
+        validated_sql: str | None = None
+        prev_error: str = ""
+        prev_reasoning: str = ""
+        for attempt in range(_MAX_RETRIES):
+            if prev_error:
+                error_section = (
+                    f"Previous attempt reasoning: {prev_reasoning}\n"
+                    f"Previous attempt failed: {prev_error}\n"
+                    "Fix the issue above."
+                )
+            else:
+                error_section = ""
+            try:
+                prompt_text = schema_ctx + error_section + question
+                input_tokens = len(_enc.encode(prompt_text))
+                logger.info("sql generation input tokens", attempt=attempt + 1, tokens=input_tokens)
+
+                result: SQLQuery = await self._chain.ainvoke({
+                    "schema": schema_ctx,
+                    "dialect": dialect,
+                    "limit": capped_limit,
+                    "error_section": error_section,
+                    "question": question,
+                })
+                sql = result.sql.strip()
+                validation_error = self._validate(sql, full_schema, capped_limit)
+                if validation_error:
+                    prev_error = validation_error
+                    prev_reasoning = result.reasoning
+                    logger.warning("sql validation failed", attempt=attempt + 1, error=validation_error)
+                    continue
+                validated_sql = self._enforce_limit(sql, capped_limit)
+                output_tokens = len(_enc.encode(result.sql)) + len(_enc.encode(result.reasoning))
+                logger.info(
+                    "sql generated",
+                    attempt=attempt + 1,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    total_tokens=input_tokens + output_tokens,
+                    reasoning=result.reasoning,
+                )
+                break
+            except Exception as e:
+                prev_error = str(e)
+                logger.warning("sql generation error", attempt=attempt + 1, error=prev_error)
+
+        if not validated_sql:
+            logger.error("sql generation failed after retries", client_id=client_id)
+            return None
+
+        # Execute on user's DB
+        creds = decrypt_credentials_dict(client.credentials)
+        with db_pipeline_service.engine_scope(client.db_type, creds) as engine:
+            rows = await asyncio.to_thread(self._run_sql, engine, validated_sql)
+
+        column_types = {
+            col["name"]: col["type"]
+            for cols in full_schema.values()
+            for col in cols
+        }
+        columns = list(rows[0].keys()) if rows else []
+
+        return QueryResult(
+            source_type="database",
+            source_id=client_id,
+            table_or_file=", ".join(table_names),
+            columns=columns,
+            rows=rows,
+            row_count=len(rows),
+            metadata={
+                "db_type": client.db_type,
+                "client_name": client.name,
+                "sql": validated_sql,
+                "column_types": {c: column_types.get(c, "unknown") for c in columns},
+            },
+        )
+
+    # ------------------------------------------------------------------
+    # Schema helpers
+    # ------------------------------------------------------------------
+
+    async def _expand_with_fk_tables(
+        self,
+        client_id: str,
+        user_id: str,
+        table_names: list[str],
+    ) -> list[str]:
+        """Expand table_names with any tables FK-referenced by the retrieved tables.
+
+        Prevents SQL generation failures when a required table (e.g. orders) wasn't
+        returned by retrieval but is referenced via FK from a table that was
+        (e.g. order_items.order_id -> orders.id).
+        """
+        if not table_names:
+            return table_names
+
+        placeholders = ", ".join(f":t{i}" for i in range(len(table_names)))
+        sql = text(f"""
+            SELECT DISTINCT lpe.cmetadata->'data'->>'foreign_key' AS fk
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'database'
+              AND lpe.cmetadata->>'database_client_id' = :client_id
+              AND lpe.cmetadata->'data'->>'table_name' IN ({placeholders})
+              AND lpe.cmetadata->'data'->>'foreign_key' IS NOT NULL
+        """)
+
+        params: dict[str, Any] = {"user_id": user_id, "client_id": client_id}
+        for i, name in enumerate(table_names):
+            params[f"t{i}"] = name
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, params)
+            rows = result.fetchall()
+
+        expanded = set(table_names)
+        for row in rows:
+            fk = row.fk  # format: "referred_table.referred_column"
+            if fk:
+                referred_table = fk.split(".")[0]
+                expanded.add(referred_table)
+
+        if expanded != set(table_names):
+            logger.info(
+                "expanded tables via FK",
+                original=sorted(table_names),
+                expanded=sorted(expanded),
+            )
+
+        return list(expanded)
+
+    async def _fetch_full_schema(
+        self,
+        client_id: str,
+        table_names: list[str],
+        user_id: str,
+    ) -> dict[str, list[dict[str, Any]]]:
+        """Fetch ALL column chunks for the given tables from PGVector.
+
+        Returns {table_name: [{"name": ..., "type": ..., "is_primary_key": ...,
+                                "foreign_key": ..., "content": ...}]}
+        """
+        placeholders = ", ".join(f":t{i}" for i in range(len(table_names)))
+        sql = text(f"""
+            SELECT lpe.cmetadata, lpe.document
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'database'
+              AND lpe.cmetadata->>'database_client_id' = :client_id
+              AND lpe.cmetadata->'data'->>'table_name' IN ({placeholders})
+            ORDER BY lpe.cmetadata->'data'->>'table_name', lpe.cmetadata->'data'->>'column_name'
+        """)
+
+        params: dict[str, Any] = {"user_id": user_id, "client_id": client_id}
+        for i, name in enumerate(table_names):
+            params[f"t{i}"] = name
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, params)
+            rows = result.fetchall()
+
+        schema: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        for row in rows:
+            data = row.cmetadata.get("data", {})
+            table = data.get("table_name")
+            if table:
+                schema[table].append({
+                    "name": data.get("column_name", ""),
+                    "type": data.get("column_type", ""),
+                    "is_primary_key": data.get("is_primary_key", False),
+                    "foreign_key": data.get("foreign_key"),
+                    "content": row.document,  # chunk text includes top values / samples
+                })
+        return dict(schema)
+
+    def _build_schema_context(self, schema: dict[str, list[dict[str, Any]]]) -> str:
+        lines: list[str] = []
+        for table, columns in schema.items():
+            lines.append(f"Table: {table}")
+            for col in columns:
+                flags = []
+                if col["is_primary_key"]:
+                    flags.append("PRIMARY KEY")
+                if col["foreign_key"]:
+                    flags.append(f"FK -> {col['foreign_key']}")
+                flag_str = f"  [{', '.join(flags)}]" if flags else ""
+                lines.append(f"  - {col['name']}  {col['type']}{flag_str}")
+                # Include sample/top-values line from chunk content if present
+                for line in col["content"].splitlines():
+                    if line.startswith(("Top values:", "Sample values:")):
+                        lines.append(f"    {line}")
+                        break
+            lines.append("")
+        return "\n".join(lines).strip()
+
+    # ------------------------------------------------------------------
+    # Guardrails
+    # ------------------------------------------------------------------
+
+    def _validate(self, sql: str, schema: dict[str, list[dict]], limit: int) -> str:
+        """Return an error string if validation fails, empty string if OK."""
+        # Layer 1: sqlglot parse + SELECT-only check
+        try:
+            parsed = sqlglot.parse_one(sql)
+        except sqlglot.errors.ParseError as e:
+            return f"SQL parse error: {e}"
+
+        if not isinstance(parsed, exp.Select):
+            return f"Only SELECT statements are allowed. Got: {type(parsed).__name__}"
+
+        # Check for DML anywhere in the AST (including writeable CTEs)
+        for node in parsed.find_all((exp.Insert, exp.Update, exp.Delete)):
+            return f"DML ({type(node).__name__}) is not allowed."
+
+        # Layer 2: schema grounding — table names
+        known_tables = {t.lower() for t in schema}
+        for tbl in parsed.find_all(exp.Table):
+            name = tbl.name.lower()
+            if name and name not in known_tables:
+                return f"Unknown table '{tbl.name}'. Only use tables from the schema."
+
+        # Layer 3: LIMIT enforcement (inject if missing — done before execution)
+        return ""
+
+    # ------------------------------------------------------------------
+    # SQL execution
+    # ------------------------------------------------------------------
+
+    def _enforce_limit(self, sql: str, limit: int) -> str:
+        """Inject or cap LIMIT using sqlglot AST manipulation."""
+        parsed = sqlglot.parse_one(sql)
+        existing = parsed.find(exp.Limit)
+        if existing:
+            current = int(existing.expression.this)
+            if current > limit:
+                existing.expression.set("this", limit)
+        else:
+            parsed = parsed.limit(limit)
+        return parsed.sql()
+
+    def _run_sql(self, engine: Any, sql: str) -> list[dict]:
+        # Ensure the user DB connection is a read-only credential — sqlglot validation alone is not sufficient.
+        with engine.connect() as conn:
+            result = conn.execute(text(sql))
+            return [dict(row) for row in result.mappings()]
+
+
+db_executor = DbExecutor()

src/query/executors/tabular.py

@@ -0,0 +1,39 @@
+"""Executor for tabular document sources (source_type="document", file_type csv/xlsx).
+
+Flow:
+  1. Group RetrievalResult chunks by document_id.
+  2. For each document: download bytes from Azure Blob -> read with pandas.
+  3. Filter DataFrame to relevant columns identified by retrieval.
+  4. Return QueryResult per document.
+"""
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src.middlewares.logging import get_logger
+from src.query.base import BaseExecutor, QueryResult
+from src.rag.base import RetrievalResult
+
+logger = get_logger("tabular_executor")
+
+_TABULAR_FILE_TYPES = ("csv", "xlsx")
+
+
+class TabularExecutor(BaseExecutor):
+    async def execute(
+        self,
+        results: list[RetrievalResult],
+        user_id: str,
+        db: AsyncSession,
+        limit: int = 100,
+    ) -> list[QueryResult]:
+        # TODO: implement
+        # 1. filter results where source_type == "document" and file_type in _TABULAR_FILE_TYPES
+        # 2. group by document_id -> list of column_names
+        # 3. per group: look up Document by document_id -> get blob_name
+        # 4. blob_storage.download_file(blob_name) -> pd.read_csv / pd.read_excel
+        # 5. df[relevant_columns].head(limit) -> rows as list[dict]
+        # 6. return QueryResult per document
+        raise NotImplementedError
+
+
+tabular_executor = TabularExecutor()

src/query/query_executor.py

@@ -0,0 +1,52 @@
+"""QueryExecutor — dispatches retrieval results to the appropriate executor by source_type."""
+
+import asyncio
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src.middlewares.logging import get_logger
+from src.query.base import QueryResult
+from src.query.executors.db_executor import db_executor
+from src.query.executors.tabular import tabular_executor
+from src.rag.base import RetrievalResult
+
+logger = get_logger("query_executor")
+
+
+class QueryExecutor:
+    async def execute(
+        self,
+        results: list[RetrievalResult],
+        user_id: str,
+        db: AsyncSession,
+        question: str,
+        limit: int = 100,
+    ) -> list[QueryResult]:
+        db_results = [r for r in results if r.source_type == "database"]
+        tabular_results = [
+            r for r in results
+            if r.source_type == "document"
+            and r.metadata.get("data", {}).get("file_type") in ("csv", "xlsx")
+        ]
+
+        async def _empty() -> list[QueryResult]:
+            return []
+
+        batches = await asyncio.gather(
+            db_executor.execute(db_results, user_id, db, question, limit) if db_results else _empty(),
+            tabular_executor.execute(tabular_results, user_id, db, question, limit) if tabular_results else _empty(),
+            return_exceptions=True,
+        )
+
+        query_results: list[QueryResult] = []
+        for batch in batches:
+            if isinstance(batch, Exception):
+                logger.error("executor failed", error=str(batch))
+                continue
+            query_results.extend(batch)
+
+        logger.info("query execution complete", total=len(query_results))
+        return query_results
+
+
+query_executor = QueryExecutor()

src/rag/base.py

@@ -0,0 +1,20 @@
+"""Shared contract for all retriever implementations."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class RetrievalResult:
+    content: str
+    metadata: dict[str, Any]
+    score: float
+    source_type: str  # "document" | "database"
+
+
+class BaseRetriever(ABC):
+    @abstractmethod
+    async def retrieve(
+        self, query: str, user_id: str, k: int = 5
+    ) -> list[RetrievalResult]: ...

src/rag/retriever.py

@@ -1,69 +1,43 @@
-"""Service for retrieving relevant documents from vector store."""
+"""Public retrieval API — thin wrapper around RetrievalRouter."""
+
+from typing import Any
 
-import hashlib
-import json
-from src.db.postgres.vector_store import get_vector_store
-from src.db.redis.connection import get_redis
 from sqlalchemy.ext.asyncio import AsyncSession
+
 from src.middlewares.logging import get_logger
-from typing import List, Dict, Any
+from src.rag.retrievers.document import document_retriever
+from src.rag.retrievers.schema import schema_retriever
+from src.rag.router import RetrievalRouter, SourceHint
 
 logger = get_logger("retriever")
 
-_RETRIEVAL_CACHE_TTL = 3600  # 1 hour
-
 
 class RetrieverService:
-    """Service for retrieving relevant documents."""
+    """Public retrieval service used by chat.py and search tools.
+
+    Delegates to RetrievalRouter which dispatches based on source_hint.
+    Returns List[Dict] to preserve backward compatibility with chat.py.
+    """
 
     def __init__(self):
-        self.vector_store = get_vector_store()
+        self._router = RetrievalRouter(
+            schema_retriever=schema_retriever,
+            document_retriever=document_retriever,
+        )
 
     async def retrieve(
         self,
         query: str,
         user_id: str,
         db: AsyncSession,
-        k: int = 5
-    ) -> List[Dict[str, Any]]:
-        """Retrieve relevant chunks for a query, scoped to the user's documents.
-
-        Returns:
-            List of dicts with keys: content, metadata
-            metadata includes: document_id, user_id, filename, chunk_index, page_label (if PDF)
-        """
+        k: int = 5,
+        source_hint: SourceHint = "both",
+    ) -> list[dict[str, Any]]:
         try:
-            redis = await get_redis()
-            query_hash = hashlib.md5(query.encode()).hexdigest()
-            cache_key = f"retrieval:{user_id}:{query_hash}:{k}"
-
-            cached = await redis.get(cache_key)
-            if cached:
-                logger.info("Returning cached retrieval results")
-                return json.loads(cached)
-
-            logger.info(f"Retrieving for user {user_id}, query: {query[:50]}...")
-
-            docs = await self.vector_store.asimilarity_search(
-                query=query,
-                k=k,
-                filter={"user_id": user_id}
-            )
-
-            results = [
-                {
-                    "content": doc.page_content,
-                    "metadata": doc.metadata,
-                }
-                for doc in docs
-            ]
-
-            logger.info(f"Retrieved {len(results)} chunks")
-            await redis.setex(cache_key, _RETRIEVAL_CACHE_TTL, json.dumps(results))
-            return results
-
+            results = await self._router.retrieve(query, user_id, source_hint, k)
+            return [{"content": r.content, "metadata": r.metadata} for r in results]
         except Exception as e:
-            logger.error("Retrieval failed", error=str(e))
+            logger.error("retrieval failed", error=str(e))
             return []
 
 

src/rag/retrievers/baseline.py

@@ -0,0 +1,70 @@
+"""Service for retrieving relevant documents from vector store."""
+
+import hashlib
+import json
+from src.db.postgres.vector_store import get_vector_store
+from src.db.redis.connection import get_redis
+from sqlalchemy.ext.asyncio import AsyncSession
+from src.middlewares.logging import get_logger
+from typing import List, Dict, Any
+
+logger = get_logger("retriever")
+
+_RETRIEVAL_CACHE_TTL = 3600  # 1 hour
+
+
+class RetrieverService:
+    """Service for retrieving relevant documents."""
+
+    def __init__(self):
+        self.vector_store = get_vector_store()
+
+    async def retrieve(
+        self,
+        query: str,
+        user_id: str,
+        db: AsyncSession,
+        k: int = 5
+    ) -> List[Dict[str, Any]]:
+        """Retrieve relevant chunks for a query, scoped to the user's documents.
+
+        Returns:
+            List of dicts with keys: content, metadata
+            metadata includes: document_id, user_id, filename, chunk_index, page_label (if PDF)
+        """
+        try:
+            redis = await get_redis()
+            query_hash = hashlib.md5(query.encode()).hexdigest()
+            cache_key = f"retrieval:{user_id}:{query_hash}:{k}"
+
+            cached = await redis.get(cache_key)
+            if cached:
+                logger.info("Returning cached retrieval results")
+                return json.loads(cached)
+
+            logger.info(f"Retrieving for user {user_id}, query: {query[:50]}...")
+
+            docs = await self.vector_store.asimilarity_search(
+                query=query,
+                k=k,
+                filter={"user_id": user_id}
+            )
+
+            results = [
+                {
+                    "content": doc.page_content,
+                    "metadata": doc.metadata,
+                }
+                for doc in docs
+            ]
+
+            logger.info(f"Retrieved {len(results)} chunks")
+            await redis.setex(cache_key, _RETRIEVAL_CACHE_TTL, json.dumps(results))
+            return results
+
+        except Exception as e:
+            logger.error("Retrieval failed", error=str(e))
+            return []
+
+
+retriever = RetrieverService()

src/rag/retrievers/document.py

@@ -0,0 +1,32 @@
+"""Document retriever — handles PDF, DOCX, TXT chunks (source_type="document", non-tabular).
+
+TEAMMATE: implement retrieve() below.
+Strategy: MMR (amax_marginal_relevance_search) + score threshold to avoid returning
+near-identical chunks from the same PDF page.
+Filter: source_type="document" AND data->>'file_type' NOT IN ('csv', 'xlsx')
+"""
+
+from src.db.postgres.vector_store import get_vector_store
+from src.middlewares.logging import get_logger
+from src.rag.base import BaseRetriever, RetrievalResult
+
+logger = get_logger("document_retriever")
+
+_SCORE_THRESHOLD = 0.45  # discard chunks with cosine distance above this
+
+
+class DocumentRetriever(BaseRetriever):
+    def __init__(self):
+        self.vector_store = get_vector_store()
+
+    async def retrieve(
+        self, query: str, user_id: str, k: int = 5
+    ) -> list[RetrievalResult]:
+        # TODO (teammate): implement MMR retrieval for prose documents
+        # Filter: {"user_id": user_id, "source_type": "document"}
+        # then post-filter to exclude file_type in ("csv", "xlsx")
+        logger.info("document retriever not yet implemented — returning empty")
+        return []
+
+
+document_retriever = DocumentRetriever()

src/rag/retrievers/schema.py

@@ -0,0 +1,349 @@
+"""Schema retriever — handles DB schemas (source_type="database") and tabular file
+columns stored as source_type="document" with file_type in ("csv","xlsx").
+
+Multiple retrieval strategies are exposed for benchmarking. The active strategy
+used by the router is `retrieve()`, which dispatches to ACTIVE_STRATEGY.
+Change ACTIVE_STRATEGY at module level to switch without touching the router.
+
+All strategies embed the query exactly once, then fan out to parallel SQL legs.
+
+Vector distance strategies:
+  dense_no_threshold  — cosine (<=>), no score floor, always returns k chunks
+  dense_dot           — inner product (<#>), equivalent to cosine for normalized embeddings
+  dense_l2            — L2/euclidean (<->), monotonic with cosine on unit-sphere vectors
+  hybrid              — RRF merge of dense + FTS (database + tabular)
+  hybrid_bm25         — RRF merge of dense + FTS (database only)
+"""
+
+import asyncio
+import time
+from typing import Literal
+
+from sqlalchemy import text
+
+from src.db.postgres.connection import _pgvector_engine
+from src.db.postgres.vector_store import get_vector_store
+from src.middlewares.logging import get_logger
+from src.rag.base import BaseRetriever, RetrievalResult
+
+logger = get_logger("schema_retriever")
+
+_TABULAR_FILE_TYPES = ("csv", "xlsx")
+
+Strategy = Literal["dense_no_threshold", "dense_dot", "dense_l2", "hybrid", "hybrid_bm25"]
+ACTIVE_STRATEGY: Strategy = "hybrid_bm25"
+
+
+class SchemaRetriever(BaseRetriever):
+    def __init__(self):
+        self.vector_store = get_vector_store()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    async def _embed_query(self, query: str) -> list[float]:
+        return await asyncio.to_thread(self.vector_store.embeddings.embed_query, query)
+
+    async def _search_db(
+        self, embedding: list[float], user_id: str, k: int, operator: str = "<=>"
+    ) -> list[RetrievalResult]:
+        """Vector search over database chunks. Accepts a pre-computed embedding."""
+        emb_str = "[" + ",".join(str(x) for x in embedding) + "]"
+
+        if operator == "<#>":
+            score_sql = f"(lpe.embedding <#> '{emb_str}'::vector) * -1"
+        elif operator == "<->":
+            score_sql = f"1.0 / (1.0 + (lpe.embedding <-> '{emb_str}'::vector))"
+        else:
+            score_sql = f"1.0 - (lpe.embedding <=> '{emb_str}'::vector)"
+
+        sql = text(f"""
+            SELECT lpe.document, lpe.cmetadata, {score_sql} AS score
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'database'
+            ORDER BY lpe.embedding {operator} '{emb_str}'::vector ASC
+            LIMIT :k
+        """)
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, {"user_id": user_id, "k": k * 4})
+            rows = result.fetchall()
+
+        return [
+            RetrievalResult(
+                content=row.document,
+                metadata=row.cmetadata,
+                score=float(row.score),
+                source_type="database",
+            )
+            for row in rows
+        ]
+
+    async def _search_tabular(
+        self, embedding: list[float], user_id: str, k: int, operator: str = "<=>"
+    ) -> list[RetrievalResult]:
+        """Vector search over tabular document chunks. Accepts a pre-computed embedding."""
+        emb_str = "[" + ",".join(str(x) for x in embedding) + "]"
+
+        if operator == "<#>":
+            score_sql = f"(lpe.embedding <#> '{emb_str}'::vector) * -1"
+        elif operator == "<->":
+            score_sql = f"1.0 / (1.0 + (lpe.embedding <-> '{emb_str}'::vector))"
+        else:
+            score_sql = f"1.0 - (lpe.embedding <=> '{emb_str}'::vector)"
+
+        sql = text(f"""
+            SELECT lpe.document, lpe.cmetadata, {score_sql} AS score
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'document'
+              AND (lpe.cmetadata->'data'->>'file_type' = 'csv'
+                OR lpe.cmetadata->'data'->>'file_type' = 'xlsx')
+            ORDER BY lpe.embedding {operator} '{emb_str}'::vector ASC
+            LIMIT :k
+        """)
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, {"user_id": user_id, "k": k * 4})
+            rows = result.fetchall()
+
+        results = []
+        for row in rows:
+            results.append(
+                RetrievalResult(
+                    content=row.document,
+                    metadata=row.cmetadata,
+                    score=float(row.score),
+                    source_type="document",
+                )
+            )
+            if len(results) >= k:
+                break
+        return results
+
+    async def _search_fts_db(self, query: str, user_id: str, k: int) -> list[RetrievalResult]:
+        """Full-text search over DB schema chunks using PostgreSQL tsvector.
+
+        Requires GIN index on langchain_pg_embedding.document (created by init_db.py).
+        """
+        sql = text("""
+            SELECT lpe.document, lpe.cmetadata,
+                   ts_rank(to_tsvector('english', lpe.document),
+                           plainto_tsquery('english', :query)) AS rank
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'database'
+              AND to_tsvector('english', lpe.document) @@ plainto_tsquery('english', :query)
+            ORDER BY rank DESC
+            LIMIT :k
+        """)
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, {"query": query, "user_id": user_id, "k": k})
+            rows = result.fetchall()
+
+        return [
+            RetrievalResult(
+                content=row.document,
+                metadata=row.cmetadata,
+                score=float(row.rank),
+                source_type="database",
+            )
+            for row in rows
+        ]
+
+    async def _search_fts_tabular(self, query: str, user_id: str, k: int) -> list[RetrievalResult]:
+        """Full-text search over tabular document chunks using PostgreSQL tsvector."""
+        sql = text("""
+            SELECT lpe.document, lpe.cmetadata,
+                   ts_rank(to_tsvector('english', lpe.document),
+                           plainto_tsquery('english', :query)) AS rank
+            FROM langchain_pg_embedding lpe
+            JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid
+            WHERE lpc.name = 'document_embeddings'
+              AND lpe.cmetadata->>'user_id' = :user_id
+              AND lpe.cmetadata->>'source_type' = 'document'
+              AND (lpe.cmetadata->'data'->>'file_type' = 'csv'
+                OR lpe.cmetadata->'data'->>'file_type' = 'xlsx')
+              AND to_tsvector('english', lpe.document) @@ plainto_tsquery('english', :query)
+            ORDER BY rank DESC
+            LIMIT :k
+        """)
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, {"query": query, "user_id": user_id, "k": k})
+            rows = result.fetchall()
+
+        return [
+            RetrievalResult(
+                content=row.document,
+                metadata=row.cmetadata,
+                score=float(row.rank),
+                source_type="document",
+            )
+            for row in rows
+        ]
+
+    def _rrf_merge(
+        self,
+        *ranked_lists: list[RetrievalResult],
+        k_rrf: int = 60,
+        top_k: int = 5,
+    ) -> list[RetrievalResult]:
+        """Reciprocal Rank Fusion — combines ranked lists using rank positions only."""
+        scores: dict[tuple, float] = {}
+        index: dict[tuple, RetrievalResult] = {}
+
+        for ranked in ranked_lists:
+            for rank, result in enumerate(ranked):
+                data = result.metadata.get("data", {})
+                key = (data.get("table_name"), data.get("column_name") or data.get("filename"))
+                scores[key] = scores.get(key, 0.0) + 1.0 / (k_rrf + rank + 1)
+                if key not in index or result.score > index[key].score:
+                    index[key] = result
+
+        def _key(r: RetrievalResult) -> tuple:
+            d = r.metadata.get("data", {})
+            return (d.get("table_name"), d.get("column_name") or d.get("filename"))
+
+        merged = sorted(index.values(), key=lambda r: scores[_key(r)], reverse=True)
+        return merged[:top_k]
+
+    def _dedup(self, results: list[RetrievalResult]) -> list[RetrievalResult]:
+        """Deduplicate by (table_name, column_name), keeping highest score per unique column."""
+        seen: dict[tuple, RetrievalResult] = {}
+        for r in results:
+            data = r.metadata.get("data", {})
+            key = (data.get("table_name"), data.get("column_name") or data.get("filename"))
+            if key not in seen or r.score > seen[key].score:
+                seen[key] = r
+        return sorted(seen.values(), key=lambda r: r.score, reverse=True)
+
+    # ------------------------------------------------------------------
+    # Named strategies — one embed call each, legs run in parallel
+    # ------------------------------------------------------------------
+
+    async def dense_no_threshold(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
+        """Cosine similarity, no score cutoff — always returns k chunks."""
+        embedding = await self._embed_query(query)
+        db_results, tabular_results = await asyncio.gather(
+            self._search_db(embedding, user_id, k),
+            self._search_tabular(embedding, user_id, k),
+        )
+        return self._dedup(db_results + tabular_results)[:k]
+
+    async def dense_dot(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
+        """Inner product similarity (<#>).
+
+        For L2-normalized embeddings (OpenAI), ranking is identical to cosine.
+        Score = raw inner product (not bounded to [0,1]).
+        """
+        embedding = await self._embed_query(query)
+        db_results, tabular_results = await asyncio.gather(
+            self._search_db(embedding, user_id, k, "<#>"),
+            self._search_tabular(embedding, user_id, k, "<#>"),
+        )
+        return self._dedup(db_results + tabular_results)[:k]
+
+    async def dense_l2(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
+        """L2 (Euclidean) distance similarity (<->).
+
+        For L2-normalized embeddings (OpenAI), ranking order matches cosine.
+        Score = 1 / (1 + l2_distance), bounded to (0, 1].
+        """
+        embedding = await self._embed_query(query)
+        db_results, tabular_results = await asyncio.gather(
+            self._search_db(embedding, user_id, k, "<->"),
+            self._search_tabular(embedding, user_id, k, "<->"),
+        )
+        return self._dedup(db_results + tabular_results)[:k]
+
+    async def hybrid(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
+        """RRF merge of dense + FTS over both database and tabular sources.
+
+        Embeds once, then runs all four legs (dense db, dense tabular, fts db,
+        fts tabular) in a single asyncio.gather.
+        """
+        embedding = await self._embed_query(query)
+        db_results, tabular_results, fts_db, fts_tabular = await asyncio.gather(
+            self._search_db(embedding, user_id, k),
+            self._search_tabular(embedding, user_id, k),
+            self._search_fts_db(query, user_id, k * 4),
+            self._search_fts_tabular(query, user_id, k * 4),
+        )
+        dense = self._dedup(db_results + tabular_results)[:k]
+        fts_all = self._dedup(fts_db + fts_tabular)
+        return self._rrf_merge(dense, fts_all, top_k=k)
+
+    async def hybrid_bm25(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
+        """RRF merge of dense + FTS (database chunks only).
+
+        Embeds once, then runs dense db, dense tabular, and fts db legs in parallel.
+        """
+        embedding = await self._embed_query(query)
+        db_results, tabular_results, fts_results = await asyncio.gather(
+            self._search_db(embedding, user_id, k),
+            self._search_tabular(embedding, user_id, k),
+            self._search_fts_db(query, user_id, k * 4),
+        )
+        dense = self._dedup(db_results + tabular_results)[:k]
+        return self._rrf_merge(dense, self._dedup(fts_results), top_k=k)
+
+    # ------------------------------------------------------------------
+    # Public interface — called by the router
+    # ------------------------------------------------------------------
+
+    async def retrieve(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
+        strategy_fn = getattr(self, ACTIVE_STRATEGY)
+        results = await strategy_fn(query, user_id, k)
+        logger.info("schema retrieval", strategy=ACTIVE_STRATEGY, count=len(results))
+        return results
+
+
+# ------------------------------------------------------------------
+# Benchmark helper — import in test scripts
+# ------------------------------------------------------------------
+
+async def benchmark(
+    query: str,
+    user_id: str,
+    k: int = 5,
+    strategies: list[Strategy] | None = None,
+) -> dict[str, dict]:
+    """Run multiple strategies against the same query and return timing + results."""
+    retriever = SchemaRetriever()
+    targets: list[Strategy] = strategies or [
+        "dense_no_threshold",
+        "dense_dot",
+        "dense_l2",
+        "hybrid",
+        "hybrid_bm25",
+    ]
+    report: dict[str, dict] = {}
+
+    for name in targets:
+        fn = getattr(retriever, name)
+        t0 = time.perf_counter()
+        chunks = await fn(query, user_id, k)
+        elapsed_ms = round((time.perf_counter() - t0) * 1000)
+
+        total_chars = sum(len(r.content) for r in chunks)
+        report[name] = {
+            "chunks": len(chunks),
+            "estimated_tokens": total_chars // 4,
+            "elapsed_ms": elapsed_ms,
+            "results": chunks,
+        }
+
+    return report
+
+
+schema_retriever = SchemaRetriever()

src/rag/router.py

@@ -0,0 +1,75 @@
+"""Routes retrieval requests to the appropriate retriever based on source_hint."""
+
+import asyncio
+import hashlib
+import json
+from typing import Literal
+
+from src.db.redis.connection import get_redis
+from src.middlewares.logging import get_logger
+from src.rag.base import BaseRetriever, RetrievalResult
+
+logger = get_logger("retrieval_router")
+
+_CACHE_TTL = 3600  # 1 hour
+SourceHint = Literal["document", "schema", "both"]
+
+
+class RetrievalRouter:
+    def __init__(
+        self,
+        schema_retriever: BaseRetriever,
+        document_retriever: BaseRetriever,
+    ):
+        self._retrievers: dict[str, BaseRetriever] = {
+            "schema": schema_retriever,
+            "document": document_retriever,
+        }
+
+    def _route(self, source_hint: SourceHint) -> list[BaseRetriever]:
+        if source_hint == "schema":
+            return [self._retrievers["schema"]]
+        if source_hint == "document":
+            return [self._retrievers["document"]]
+        return list(self._retrievers.values())
+
+    async def retrieve(
+        self,
+        query: str,
+        user_id: str,
+        source_hint: SourceHint = "both",
+        k: int = 10,
+    ) -> list[RetrievalResult]:
+        redis = await get_redis()
+        query_hash = hashlib.md5(query.encode()).hexdigest()
+        cache_key = f"retrieval:{user_id}:{source_hint}:{query_hash}:{k}"
+
+        cached = await redis.get(cache_key)
+        if cached:
+            logger.info("returning cached retrieval results", source_hint=source_hint)
+            raw = json.loads(cached)
+            return [RetrievalResult(**r) for r in raw]
+
+        retrievers = self._route(source_hint)
+        batches = await asyncio.gather(
+            *[r.retrieve(query, user_id, k) for r in retrievers],
+            return_exceptions=True,
+        )
+
+        results: list[RetrievalResult] = []
+        for batch in batches:
+            if isinstance(batch, Exception):
+                logger.error("retriever failed", error=str(batch))
+                continue
+            results.extend(batch)
+
+        results.sort(key=lambda r: r.score, reverse=True)
+        results = results[:k]
+
+        logger.info("retrieved chunks", count=len(results), source_hint=source_hint)
+        await redis.setex(
+            cache_key,
+            _CACHE_TTL,
+            json.dumps([vars(r) for r in results]),
+        )
+        return results

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()

uv.lock

@@ -66,6 +66,7 @@ dependencies = [
     { name = "spacy" },
     { name = "sqlalchemy", extra = ["asyncio"] },
     { name = "sqlalchemy-bigquery" },
+    { name = "sqlglot" },
     { name = "sse-starlette" },
     { name = "starlette" },
     { name = "structlog" },
@@ -149,6 +150,7 @@ requires-dist = [
     { name = "spacy", specifier = "==3.8.3" },
     { name = "sqlalchemy", extras = ["asyncio"], specifier = "==2.0.36" },
     { name = "sqlalchemy-bigquery", specifier = ">=1.11.0" },
+    { name = "sqlglot", specifier = ">=25.0.0" },
     { name = "sse-starlette", specifier = "==2.1.3" },
     { name = "starlette", specifier = "==0.41.3" },
     { name = "structlog", specifier = "==24.4.0" },
@@ -3221,6 +3223,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c0/87/11e6de00ef7949bb8ea06b55304a1a4911c329fdf0d9882b464db240c2c5/sqlalchemy_bigquery-1.16.0-py3-none-any.whl", hash = "sha256:0fe7634cd954f3e74f5e2db6d159f9e5ee87a47fbe8d52eac3cd3bb3dadb3a77", size = 40615, upload-time = "2025-11-06T01:35:39.358Z" },
 ]
 
+[[package]]
+name = "sqlglot"
+version = "30.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/3c/66/6ece15f197874e56c76e1d0269cebf284ba992a80dfadca9d1972fdf7edf/sqlglot-30.6.0.tar.gz", hash = "sha256:246d34d39927422a50a3fa155f37b2f6346fba85f1a755b13c941eb32ef93361", size = 5835307, upload-time = "2026-04-20T20:11:08.164Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/dc/e7/64fe971cbca33a0446b06f4a5ff8e3fa4a1dbd0a039ceabcc3e6cf4087a9/sqlglot-30.6.0-py3-none-any.whl", hash = "sha256:e005fc2f47994f90d7d8df341f1cbe937518497b0b7b1507d4c03c4c9dfd2778", size = 673920, upload-time = "2026-04-20T20:11:05.758Z" },
+]
+
 [[package]]
 name = "srsly"
 version = "2.5.3"