[KM-438][KM-439] framework for knowledge retriever

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

@@ -192,6 +192,7 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
                         query=search_query,
                         user_id=request.user_id,
                         db=db,
+                        source_hint=intent_result.get("source_hint", "both"),
                     )
                 else:
                     raw_results = await retrieval_task

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/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,86 @@
+"""Schema retriever — handles DB schemas (source_type="database") and tabular file
+columns stored as source_type="document" with file_type in ("csv","xlsx").
+
+Strategy: similarity search with score threshold on two metadata shapes,
+run in parallel, merged and re-ranked by score.
+"""
+
+import asyncio
+
+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")
+
+_SCORE_THRESHOLD = 0.45  # cosine distance — discard above this value
+_TABULAR_FILE_TYPES = ("csv", "xlsx")
+
+
+class SchemaRetriever(BaseRetriever):
+    def __init__(self):
+        self.vector_store = get_vector_store()
+
+    async def _search_db(self, query: str, user_id: str, k: int) -> list[RetrievalResult]:
+        """Retrieve DB schema chunks (source_type="database")."""
+        docs_with_scores = await self.vector_store.asimilarity_search_with_score(
+            query=query,
+            k=k,
+            filter={"user_id": user_id, "source_type": "database"},
+        )
+        results = []
+        for doc, distance in docs_with_scores:
+            if distance <= _SCORE_THRESHOLD:
+                results.append(
+                    RetrievalResult(
+                        content=doc.page_content,
+                        metadata=doc.metadata,
+                        score=1.0 - distance,
+                        source_type="database",
+                    )
+                )
+        return results
+
+    async def _search_tabular(self, query: str, user_id: str, k: int) -> list[RetrievalResult]:
+        """Retrieve CSV/XLSX column chunks (source_type="document", file_type=csv|xlsx)."""
+        results = []
+        for file_type in _TABULAR_FILE_TYPES:
+            docs_with_scores = await self.vector_store.asimilarity_search_with_score(
+                query=query,
+                k=k,
+                filter={
+                    "user_id": user_id,
+                    "source_type": "document",
+                    "data": {"file_type": file_type},
+                },
+            )
+            for doc, distance in docs_with_scores:
+                if distance <= _SCORE_THRESHOLD:
+                    results.append(
+                        RetrievalResult(
+                            content=doc.page_content,
+                            metadata=doc.metadata,
+                            score=1.0 - distance,
+                            source_type="document",
+                        )
+                    )
+        return results
+
+    async def retrieve(
+        self, query: str, user_id: str, k: int = 5
+    ) -> list[RetrievalResult]:
+        db_results, tabular_results = await asyncio.gather(
+            self._search_db(query, user_id, k),
+            self._search_tabular(query, user_id, k),
+        )
+        combined = db_results + tabular_results
+        combined.sort(key=lambda r: r.score, reverse=True)
+        logger.info(
+            "schema retrieval",
+            db_chunks=len(db_results),
+            tabular_chunks=len(tabular_results),
+        )
+        return combined[:k]
+
+
+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 = 5,
+    ) -> 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