File size: 1,505 Bytes
52999bc 027123c 52999bc 027123c c93ec90 52999bc 027123c 52999bc c93ec90 52999bc 027123c 52999bc 027123c 52999bc c93ec90 027123c c93ec90 027123c 52999bc 027123c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """Public retrieval API — thin wrapper around RetrievalRouter."""
from sqlalchemy.ext.asyncio import AsyncSession
from src.middlewares.logging import get_logger
from src.rag.base import RetrievalResult
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")
class RetrieverService:
"""Public retrieval service used by chat.py and search tools.
Delegates to RetrievalRouter which dispatches based on source_hint.
Returns RetrievalResult objects directly so downstream consumers
(db_executor, tabular_executor) can be fed without lossy dict
conversion. The `db` parameter is accepted for call-site compatibility
but currently unused — retrieval reads PGVector via _pgvector_engine
inside each retriever.
"""
def __init__(self):
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,
source_hint: SourceHint = "both",
) -> list[RetrievalResult]:
try:
return await self._router.retrieve(query, user_id, source_hint, k)
except Exception as e:
logger.error("retrieval failed", error=str(e))
return []
retriever = RetrieverService()
|