gqt
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
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
|