[KM-533] now also retrieves table level chunk

src/rag/retrievers/schema.py

@@ -1,9 +1,15 @@
 """Schema retriever — handles DB schemas (source_type="database") and tabular file
 columns stored as source_type="document" with file_type in ("csv","xlsx").
 
-Strategy: hybrid_bm25 — RRF merge of dense cosine search (DB + tabular) and
-PostgreSQL full-text search (DB only). Embeds the query once, fans out the
-three legs in parallel.
+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.
+
+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
+don't individually win on similarity, the table chunk can still pull the table
+into the hit set, where db_executor's downstream full-schema fetch picks up
+the per-column detail.
 
 FTS requires a GIN index on langchain_pg_embedding.document (created by init_db.py).
 """
@@ -20,6 +26,7 @@ from src.rag.base import BaseRetriever, RetrievalResult
 logger = get_logger("schema_retriever")
 
 _TABULAR_FILE_TYPES = ("csv", "xlsx")
+_TABLE_CHUNK_K_MULTIPLIER = 2  # how many table chunks to pull before RRF
 
 
 class SchemaRetriever(BaseRetriever):
@@ -66,6 +73,46 @@ class SchemaRetriever(BaseRetriever):
             for row in rows
         ]
 
+    async def _search_db_tables(
+        self, embedding: list[float], user_id: str, k: int
+    ) -> list[RetrievalResult]:
+        """Cosine vector search over database TABLE-level chunks.
+
+        Recall channel for multi-table questions. The chunk's content is
+        discarded downstream — db_executor only consumes its `data.table_name`
+        to seed full-schema fetch.
+        """
+        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' = 'database'
+              AND lpe.cmetadata->>'chunk_level' = 'table'
+            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 * _TABLE_CHUNK_K_MULTIPLIER}
+            )
+            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
     ) -> list[RetrievalResult]:
@@ -171,16 +218,24 @@ class SchemaRetriever(BaseRetriever):
     # ------------------------------------------------------------------
 
     async def retrieve(self, query: str, user_id: str, k: int = 5) -> list[RetrievalResult]:
-        """RRF merge of dense (DB + tabular) and FTS (DB only)."""
+        """RRF merge of dense (DB columns + DB tables + tabular) and FTS (DB cols only)."""
         embedding = await self._embed_query(query)
-        db_results, tabular_results, fts_results = await asyncio.gather(
+        db_col_results, db_tbl_results, tabular_results, fts_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),
         )
-        dense = self._dedup(db_results + tabular_results)[:k]
+        dense = self._dedup(db_col_results + db_tbl_results + tabular_results)[:k]
         results = self._rrf_merge(dense, self._dedup(fts_results), top_k=k)
-        logger.info("schema retrieval", count=len(results))
+        logger.info(
+            "schema retrieval",
+            count=len(results),
+            db_cols=len(db_col_results),
+            db_tables=len(db_tbl_results),
+            tabular=len(tabular_results),
+            fts=len(fts_results),
+        )
         return results