[NOTICKET][db] add sheet-level retrieval and focus LLM schema context to retrieved columns

src/query/executors/tabular.py

@@ -252,6 +252,10 @@ class TabularExecutor(BaseExecutor):
         ) -> QueryResult | None:
             try:
                 df = await download_parquet(user_id, doc_id, sheet_name)
+                if info["columns"]:
+                    valid_cols = [c for c in info["columns"] if c in df.columns]
+                    if valid_cols:
+                        df = df[valid_cols]
                 df_result = await self._query_with_agent(df, question, limit)
 
                 table_label = info["filename"]

src/rag/retrievers/schema.py

@@ -2,8 +2,8 @@
 columns stored as source_type="document" with file_type in ("csv","xlsx").
 
 Strategy: hybrid_bm25 — RRF merge of dense cosine search (DB columns + DB tables
-+ tabular) and PostgreSQL full-text search (DB columns only). Embeds the query
-once, fans out four legs in parallel.
++ tabular columns + tabular sheets) and PostgreSQL full-text search (DB columns only).
+Embeds the query once, fans out five legs in parallel.
 
 The DB-tables leg surfaces table-level summary chunks (chunk_level='table') as
 a recall signal for multi-table questions: when a relevant table's columns
@@ -127,6 +127,7 @@ class SchemaRetriever(BaseRetriever):
             WHERE lpc.name = 'document_embeddings'
               AND lpe.cmetadata->>'user_id' = :user_id
               AND lpe.cmetadata->>'source_type' = 'document'
+              AND lpe.cmetadata->>'chunk_level' = 'column'
               AND (lpe.cmetadata->'data'->>'file_type' = 'csv'
                 OR lpe.cmetadata->'data'->>'file_type' = 'xlsx')
             ORDER BY lpe.embedding <=> '{emb_str}'::vector ASC
@@ -147,6 +148,41 @@ class SchemaRetriever(BaseRetriever):
             for row in rows
         ]
 
+    async def _search_tabular_sheets(
+        self, embedding: list[float], user_id: str, k: int
+    ) -> list[RetrievalResult]:
+        """Leg 5: sheet-level summary chunks from CSV/XLSX files."""
+        emb_str = "[" + ",".join(str(x) for x in embedding) + "]"
+
+        sql = text(f"""
+            SELECT lpe.document, lpe.cmetadata,
+                   1.0 - (lpe.embedding <=> '{emb_str}'::vector) 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->>'chunk_level' = 'sheet'
+              AND (lpe.cmetadata->'data'->>'file_type' = 'csv'
+                OR lpe.cmetadata->'data'->>'file_type' = 'xlsx')
+            ORDER BY lpe.embedding <=> '{emb_str}'::vector ASC
+            LIMIT :k
+        """)
+
+        async with _pgvector_engine.connect() as conn:
+            result = await conn.execute(sql, {"user_id": user_id, "k": k})
+            rows = result.fetchall()
+
+        return [
+            RetrievalResult(
+                content=row.document,
+                metadata=row.cmetadata,
+                score=float(row.score),
+                source_type="document",
+            )
+            for row in rows
+        ]
+
     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."""
         sql = text("""
@@ -182,9 +218,10 @@ class SchemaRetriever(BaseRetriever):
     def _chunk_key(r: RetrievalResult) -> tuple:
         """Stable identity for dedup/RRF.
 
-        Includes filename and sheet_name so that tabular column chunks with
-        the same column name across different files (e.g. `id` in two CSVs)
-        and future sheet-level chunks across XLSX sheets don't collide.
+        Includes filename, sheet_name, and chunk_level so that column chunks
+        and sheet chunks for the same file/sheet don't collide, and column
+        chunks with the same name across different files (e.g. `id` in two CSVs)
+        are kept distinct.
         """
         d = r.metadata.get("data", {})
         return (
@@ -192,6 +229,7 @@ class SchemaRetriever(BaseRetriever):
             d.get("column_name"),
             d.get("filename"),
             d.get("sheet_name"),
+            r.metadata.get("chunk_level"),
         )
 
     def _dedup(self, results: list[RetrievalResult]) -> list[RetrievalResult]:
@@ -291,15 +329,16 @@ class SchemaRetriever(BaseRetriever):
         no table-level chunks.
         """
         embedding = await self._embed_query(query)
-        db_col_results, db_tbl_results, tabular_results, fts_results = await asyncio.gather(
+        db_col_results, db_tbl_results, tabular_results, fts_results, sheet_results = await asyncio.gather(
             self._search_db(embedding, user_id, k),
             self._search_db_tables(embedding, user_id, k),
             self._search_tabular(embedding, user_id, k),
             self._search_fts_db(query, user_id, k * 4),
+            self._search_tabular_sheets(embedding, user_id, k),
         )
 
         db_ranked = self._rank_db_tables(db_tbl_results, db_col_results, fts_results, top_k=k)
-        tabular_final = self._dedup(tabular_results)[:k]
+        tabular_final = self._dedup(tabular_results + sheet_results)[:k]
 
         results = db_ranked + tabular_final
         logger.info(
@@ -309,6 +348,7 @@ class SchemaRetriever(BaseRetriever):
             db_cols=len(db_col_results),
             db_tables=len(db_tbl_results),
             tabular=len(tabular_results),
+            tabular_sheets=len(sheet_results),
             fts=len(fts_results),
         )
         return results

src/rag/router.py

@@ -25,7 +25,7 @@ SourceHint = Literal["document", "schema", "both"]
 
 def _result_dedup_key(r: RetrievalResult) -> tuple:
     """Cross-retriever dedup key — distinguishes DB columns vs DB tables vs
-    tabular columns vs prose chunks vs sheet-level (future)."""
+    tabular columns vs prose chunks vs sheet-level chunks."""
     data = r.metadata.get("data", {})
     return (
         r.source_type,
@@ -34,6 +34,7 @@ def _result_dedup_key(r: RetrievalResult) -> tuple:
         data.get("filename"),
         data.get("sheet_name"),
         data.get("chunk_index"),  # disambiguates multiple prose chunks per doc
+        r.metadata.get("chunk_level"),  # distinguishes sheet vs column chunks
     )