feat: add durable document indexing pipeline

backend/src/kbqa/rag/retriever.py

@@ -0,0 +1,56 @@
+import asyncio
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import async_sessionmaker
+
+from kbqa.documents.models import Chunk, Document
+from kbqa.rag.embeddings import EmbeddingProvider
+from kbqa.rag.types import SourceCandidate
+from kbqa.rag.vector_store import MilvusVectorStore
+
+
+class KnowledgeRetriever:
+    def __init__(
+        self,
+        embeddings: EmbeddingProvider,
+        vector_store: MilvusVectorStore,
+        session_factory: async_sessionmaker,
+    ) -> None:
+        self.embeddings = embeddings
+        self.vector_store = vector_store
+        self.session_factory = session_factory
+
+    async def retrieve(self, query: str, top_k: int) -> list[SourceCandidate]:
+        vector = await self.embeddings.embed_query(query)
+        candidate_limit = max(top_k * 3, top_k)
+        hits = await asyncio.to_thread(self.vector_store.search, vector, candidate_limit)
+        if not hits:
+            return []
+        ids = [hit.chunk_id for hit in hits]
+        async with self.session_factory() as session:
+            rows = await session.execute(
+                select(Chunk, Document)
+                .join(Document, Document.id == Chunk.document_id)
+                .where(Chunk.id.in_(ids), Document.status == "indexed")
+            )
+            lookup = {chunk.id: (chunk, document) for chunk, document in rows}
+        sources: list[SourceCandidate] = []
+        for hit in hits:
+            record = lookup.get(hit.chunk_id)
+            if record is None:
+                continue
+            chunk, document = record
+            sources.append(
+                SourceCandidate(
+                    chunk_id=chunk.id,
+                    document_id=document.id,
+                    filename=document.filename,
+                    order_index=chunk.order_index,
+                    page_number=chunk.page_number,
+                    content=chunk.content,
+                    score=hit.score,
+                )
+            )
+            if len(sources) == top_k:
+                break
+        return sources