backend/services/retrieval.py

"""Retrieval service wrapping FAISS hybrid search."""

from typing import Dict, List, Optional

from config import settings
from db.faiss_client import FaissDB
from models.embedder import MiniCPMEmbedder


class RetrievalService:
    def __init__(self, embedder: MiniCPMEmbedder, db: FaissDB):
        self.embedder = embedder
        self.db = db

    def search(
        self,
        query: str,
        top_k: Optional[int] = None,
        document_ids: Optional[List[str]] = None,
    ) -> List[Dict]:
        query_vector = self.embedder.embed_query(query)
        return self.db.hybrid_search(
            query_vector=query_vector,
            query_text=query,
            top_k=top_k or settings.TOP_K,
            document_ids=document_ids,
        )