File size: 2,107 Bytes
bef5e76 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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()
|