aramesh-backend/chatbot/retrieval.py

"""بازیابی ترکیبی — پورت خط لوله v13 با سه تغییر برای این سایت:

۱) پایگاه sparse به‌جای SPLADE (انگلیسی‌محور و بی‌فایده روی فارسی) از BM25
   با نرمال‌سازی فارسی استفاده می‌کند.
۲) جست‌وجوی معنایی به‌جای FAISS با pgvector و فاصله کسینوسی در SQL انجام
   می‌شود (همگام‌سازی افزایشی، فیلترپذیر، بدون فایل ایندکس جدا).
۳) مسیریابی alpha با LLM حذف شده و وزن ثابت تنظیم‌پذیر است — برای پیکرهٔ
   کوچک سایت، هزینهٔ یک فراخوانی LLM اضافی به‌ازای هر پیام نمی‌ارزد.
بقیهٔ منطق (نرمال‌سازی نمره، فیوژن وزن‌دار، برش خوشه‌بندی، ری‌رنک) عین v13 است.
"""

import logging
import time

from . import vectorstore
from .config import settings
from .indexing import INDEX
from .persian import tokenize

log = logging.getLogger(__name__)


def dense_retrieve(query: str, top_k: int) -> list[tuple[dict, float]]:
    """جست‌وجوی معنایی pgvector — خروجی: (سند، شباهت کسینوسی)."""
    return vectorstore.dense_search(query, top_k)


def keyword_retrieve(query: str, top_k: int) -> list[tuple[dict, float]]:
    """جست‌وجوی کلیدواژه BM25 روی متن نرمال‌شده فارسی."""
    if INDEX.bm25 is None:
        return []
    try:
        tokens = tokenize(query)
        if not tokens:
            return []
        scores = INDEX.bm25.get_scores(tokens)
        ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
        return [(INDEX.docs[i], float(scores[i])) for i in ranked if scores[i] > 0]
    except Exception as e:
        log.error("BM25 retrieval failed: %s", e)
        return []


def _minmax(scores: list[float]) -> list[float]:
    if not scores:
        return []
    lo, hi = min(scores), max(scores)
    if hi == lo:
        return [0.5] * len(scores)
    return [(s - lo) / (hi - lo) for s in scores]


def hybrid_retrieve(query: str, alpha: float | None = None, top_k: int | None = None) -> list[tuple[dict, float]]:
    """ترکیب وزن‌دار نمره معنایی و کلیدواژه؛ خروجی مرتب‌شده بر اساس نمره ترکیبی."""
    alpha = settings.alpha_dense if alpha is None else alpha
    top_k = top_k or settings.retrieve_k

    dense = dense_retrieve(query, top_k)
    keyword = keyword_retrieve(query, top_k)
    log.info("Retrieval '%s': dense=%d, keyword=%d", query[:60], len(dense), len(keyword))
    if not dense and not keyword:
        return []

    dense_norm = _minmax([s for _, s in dense])
    keyword_norm = _minmax([s for _, s in keyword])

    pool: dict = {}  # doc_id → {doc, dense_max, kw_max}
    for (doc, _), norm in zip(dense, dense_norm):
        entry = pool.setdefault(doc["doc_id"], {"doc": doc, "dense_max": 0.0, "kw_max": 0.0})
        entry["dense_max"] = max(entry["dense_max"], norm)
    for (doc, _), norm in zip(keyword, keyword_norm):
        entry = pool.setdefault(doc["doc_id"], {"doc": doc, "dense_max": 0.0, "kw_max": 0.0})
        entry["kw_max"] = max(entry["kw_max"], norm)

    combined = [
        (entry["doc"], alpha * entry["dense_max"] + (1 - alpha) * entry["kw_max"])
        for entry in pool.values()
    ]
    combined.sort(key=lambda x: x[1], reverse=True)
    return combined


def cluster_cutoff(scored: list[tuple[dict, float]], silhouette_threshold: float = 0.55) -> list[tuple[dict, float]]:
    """برش خوشه‌ای v13 — مرز طبیعی کیفیت به‌جای top-k دلخواه.

    روی نمره‌های ۱بعدی K-Means با K=2 می‌زند؛ اگر جدایی خوشه‌ها معنادار
    باشد (silhouette ≥ 0.55) فقط خوشهٔ بالا را نگه می‌دارد.
    """
    import numpy as np

    n = len(scored)
    if n < 4:
        return scored
    try:
        from sklearn.cluster import KMeans
        from sklearn.metrics import silhouette_score

        scores = np.array([s for _, s in scored], dtype=float).reshape(-1, 1)
        kmeans = KMeans(n_clusters=2, random_state=42, n_init=10)
        labels = kmeans.fit_predict(scores)
        sil = silhouette_score(scores, labels)
        if sil < silhouette_threshold:
            log.info("Clustering: silhouette=%.3f < %.2f — keeping all %d docs", sil, silhouette_threshold, n)
            return scored
        centers = kmeans.cluster_centers_.flatten()
        good_label = int(np.argmax(centers))
        cutoff = int(np.sum(labels == good_label))
        cutoff = max(cutoff, 1)
        log.info("Clustering: silhouette=%.3f — cutting at %d/%d docs", sil, cutoff, n)
        return scored[:cutoff]
    except Exception as e:
        log.warning("Clustering failed (%s) — keeping all docs", e)
        return scored


def rerank(query: str, scored: list[tuple[dict, float]], top_k: int) -> list[tuple[dict, float]]:
    """ری‌رنک کراس‌انکودر BGE — مرتب‌سازی نهایی بر اساس شباهت واقعی پرسش/سند."""
    if not scored:
        return []
    if not settings.reranker_enabled:
        return scored[:top_k]

    try:
        from .reranker import get_reranker

        start = time.time()
        pairs = [[query, doc["text"]] for doc, _ in scored]
        scores = get_reranker().score(pairs)
        rescored = sorted(
            ((doc, float(s)) for (doc, _), s in zip(scored, scores)),
            key=lambda x: x[1],
            reverse=True,
        )
        log.info("Reranked %d docs in %.2fs — top scores: %s",
                 len(scored), time.time() - start,
                 [f"{s:.2f}" for _, s in rescored[:3]])
        return rescored[:top_k]
    except Exception as e:
        log.error("Reranking failed (%s) — falling back to fusion order", e)
        return scored[:top_k]


def retrieve_final(query: str, query_variants: list[str] | None = None) -> list[dict]:
    """خط کامل بازیابی برای یک پرسش: هر نوع → فیوژن → برش خوشه → ری‌رنک → K نهایی."""
    variants = query_variants or [query]

    pool: dict = {}
    for variant in variants:
        for doc, score in hybrid_retrieve(variant):
            entry = pool.setdefault(doc["doc_id"], {"doc": doc, "score": 0.0})
            entry["score"] = max(entry["score"], score)

    combined = sorted(((e["doc"], e["score"]) for e in pool.values()),
                      key=lambda x: x[1], reverse=True)
    if not combined:
        return []

    filtered = cluster_cutoff(combined)
    # ری‌رنک روی نتیجه‌های برش‌خورده (سقف retrieve_k تا تأخیر کنترل شود)
    rerank_candidates = filtered[: settings.retrieve_k]
    final = rerank(query, rerank_candidates, settings.best_of_k)

    log.info("Final retrieval: %d docs (from %d candidates, %d variants)",
             len(final), len(combined), len(variants))
    return [doc for doc, _ in final]