Enterprise AI Engineering & Production Architecture

from typing import List, Dict, Any
from pydantic import BaseModel, Field
import asyncpg
from sentence_transformers import CrossEncoder

class SearchResult(BaseModel):
    chunk_id: str
    document_title: str
    content: str
    dense_score: float
    sparse_score: float
    rerank_score: float = Field(default=0.0)

class HybridRAGRetriever:
    def __init__(self, db_pool: asyncpg.Pool, reranker_model: str = "BAAI/bge-reranker-large"):
        self.db = db_pool
        self.reranker = CrossEncoder(reranker_model)

    async def retrieve(self, query: str, query_embedding: List[float], tenant_id: str, top_k: int = 5) -> List[SearchResult]:
        # Hybrid query combining pgvector HNSW cosine similarity + full-text search (ts_rank)
        sql = """
        WITH dense_search AS (
            SELECT id, document_title, content, 1 - (embedding <=> $1::vector) AS dense_sim
            FROM knowledge_embeddings
            WHERE tenant_id = $2
            ORDER BY embedding <=> $1::vector ASC
            LIMIT 25
        ),
        sparse_search AS (
            SELECT id, document_title, content, ts_rank(search_vector, plainto_tsquery('english', $3)) AS text_rank
            FROM knowledge_embeddings
            WHERE tenant_id = $2 AND search_vector @@ plainto_tsquery('english', $3)
            ORDER BY text_rank DESC
            LIMIT 25
        )
        SELECT 
            COALESCE(d.id, s.id) as chunk_id,
            COALESCE(d.document_title, s.document_title) as document_title,
            COALESCE(d.content, s.content) as content,
            COALESCE(d.dense_sim, 0.0) as dense_score,
            COALESCE(s.text_rank, 0.0) as sparse_score
        FROM dense_search d
        FULL OUTER JOIN sparse_search s ON d.id = s.id;
        """
        async with self.db.acquire() as conn:
            rows = await conn.fetch(sql, query_embedding, tenant_id, query)

        candidates = [SearchResult(**dict(row)) for row in rows]
        if not candidates:
            return []

        # Cross-encoder semantic reranking over candidate pool
        pairs = [[query, c.content] for c in candidates]
        scores = self.reranker.predict(pairs)
        for i, score in enumerate(scores):
            candidates[i].rerank_score = float(score)

        candidates.sort(key=lambda x: x.rerank_score, reverse=True)
        return candidates[:top_k]