25 lines
1.1 KiB
Python
25 lines
1.1 KiB
Python
|
"""Implement infrastructure support for dense retriever."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from app.domain.retrieval import EmbeddingProvider, RetrievalQuery, Retriever, RetrievedChunk, VectorIndex
|
||
|
|
# Keep adapter behavior explicit so integration details remain easy to audit.
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
class DenseRetriever(Retriever):
|
||
|
|
"""Provide the Dense Retriever retriever."""
|
||
|
|
def __init__(self, *, embedding_provider: EmbeddingProvider, vector_index: VectorIndex) -> None:
|
||
|
|
"""Initialize the Dense Retriever instance."""
|
||
|
|
self.embedding_provider = embedding_provider
|
||
|
|
self.vector_index = vector_index
|
||
|
|
|
||
|
|
def retrieve(self, query: RetrievalQuery) -> list[RetrievedChunk]:
|
||
|
|
"""Handle retrieve for the Dense Retriever instance."""
|
||
|
|
query_vector = self.embedding_provider.embed_query(query.query)
|
||
|
|
return self.vector_index.search(query_vector, query.top_k, query.filters)
|
||
|
|
|
||
|
|
def search(self, query: str, top_k: int, filters: str | None = None) -> list[RetrievedChunk]:
|
||
|
|
"""Handle search for the Dense Retriever instance."""
|
||
|
|
return self.retrieve(RetrievalQuery(query=query, top_k=top_k, filters=filters))
|