MathApp/rag/services/retrieval.py

"""
مرحله‌ی retrieval: پرسشِ کاربر → embedding → جستجوی Qdrant → چند خلاصه‌ی
مرتبط به‌عنوانِ context.

از UnexpectedResponseِ Qdrant (مثلاً collection موجود نیست چون این درس
هنوز آپلود نشده) به‌جای crash، لیستِ خالی برمی‌گردونیم تا لایه‌ی بالا
بتونه پاسخِ «این درس هنوز پوشش داده نشده» بده.
"""
import logging

from .config import build_collection_name, resolve_grade_title, retrieval_top_k
from .embedding import embed_text
from .known_collections import (
    KNOWN_COLLECTIONS, filter_key_for_subject, is_math_subject, resolve_math_chapter,
)
from .qdrant_store import search

logger = logging.getLogger(__name__)


def resolve_target(subject: str, grade_title: str, grade_id, lesson_number):
    """
    (collection, filter_key, filter_value) — مقصدِ دقیقِ retrieval در Qdrant.

    ریاضی با کلیدِ payloadِ `chapter` (رشته‌ی کاملِ فصل) فیلتر می‌شه و شماره‌ی
    ساده‌ی فصل به همون رشته تبدیل می‌شه؛ بقیه با `lesson_number` (که در Qdrant
    عددِ int ذخیره شده). فقط برای دیباگ/نمایشِ داخلی هم استفاده می‌شه.
    """
    grade_title = grade_title or resolve_grade_title(grade_id)
    if not subject or not grade_title:
        return None, None, None

    collection = build_collection_name(subject, grade_title)
    if collection not in KNOWN_COLLECTIONS:
        # subject و grade می‌توانند از بدنه‌ی درخواست بیایند؛ بی‌این چک هر کاربری از هر collectionِ
        # Qdrantِ مشترک (از جمله مالِ تیمِ دیگر) متن می‌خواند و در usedChunks پس می‌گرفت.
        # نتیجه همان رفتارِ collectionِ ناموجود است: جوابِ بی‌context.
        logger.warning('rag.retrieve: collection %r در KNOWN_COLLECTIONS نیست، جستجو نشد', collection)
        return None, None, None
    filter_key = filter_key_for_subject(subject)
    if is_math_subject(subject):
        filter_value = resolve_math_chapter(subject, grade_title, lesson_number)
    else:
        filter_value = lesson_number
    return collection, filter_key, filter_value


def retrieve_context(question: str, *, subject: str, grade_id=None, grade_title: str = '',
                     lesson_number=None, limit: int | None = None) -> tuple[list[dict], dict]:
    """
    خروجی: (chunks, meta)
      chunks = [{title, summary, score}, ...]  (ممکنه خالی باشه، بدونِ خطا)
      meta   = {collection, filterKey, filterValue, gradeTitle, chunkCount}
               فقط برای نمایشِ داخلی/دیباگ (نه برای کاربرِ نهایی).
    """
    grade_title = grade_title or resolve_grade_title(grade_id)
    collection, filter_key, filter_value = resolve_target(subject, grade_title, grade_id, lesson_number)
    meta = {
        'collection': collection,
        'filterKey': filter_key if filter_value is not None else None,
        'filterValue': filter_value,
        'gradeTitle': grade_title,
        'chunkCount': 0,
    }

    if collection is None:
        logger.info('rag.retrieve: subject/grade ناقص (subject=%r grade=%r) — رد شد', subject, grade_title)
        return [], meta

    limit = limit or retrieval_top_k()

    try:
        query_vector = embed_text(question)
    except Exception:
        logger.exception('rag.retrieve: embedding شکست خورد')
        raise

    try:
        from qdrant_client.http.exceptions import UnexpectedResponse
    except ModuleNotFoundError:
        UnexpectedResponse = ()  # type: ignore

    try:
        hits = search(
            collection,
            query_vector,
            lesson_filter_value=filter_value,
            lesson_filter_key=filter_key if filter_value is not None else None,
            limit=limit,
        )
    except Exception as exc:  # noqa: BLE001
        if UnexpectedResponse and isinstance(exc, UnexpectedResponse) and getattr(exc, 'status_code', None) == 404:
            logger.info('rag.retrieve: collection %r موجود نیست', collection)
            return [], meta
        logger.exception('rag.retrieve: جستجوی Qdrant در collection %r شکست خورد', collection)
        raise

    context = []
    for hit in hits:
        payload = hit['payload']
        summary = _payload_text(payload)
        if not summary:
            continue
        context.append({
            'title': _payload_title(payload),
            'summary': summary,
            'score': hit.get('score'),
        })
    meta['chunkCount'] = len(context)
    return context, meta


def _payload_text(payload: dict) -> str:
    """
    دو شکلِ payload در دادهٔ تیمِ دیگه هست:
    - درس‌های متنی (علوم/اجتماعی): کلیدِ `summary`
    - ریاضی: کلیدِ `generalSolution` (+ problem/solution) و بدونِ `summary`

    ⚠️ راه‌حلی که `rag_verify_solutions` با آن مخالفت کرده (`disagrees`) به مدل
    نمی‌رسد؛ فقط روش و صورتِ مسئله می‌رود. پیش‌تر همان راه‌حلِ غلط (مثلاً «در
    [−π, π] دو نوسانِ کامل») به‌عنوانِ «متنِ درس» به دستیار داده می‌شد و دستیار
    تکرارش می‌کرد — حدودِ ۱۳٪ تکه‌های کتاب‌های دهم تا دوازدهم این‌طورند.
    """
    if payload.get('summary'):
        return payload['summary']
    solution = payload.get('solution', '')
    if payload.get('solution_verdict') == 'disagrees':
        solution = ''
    parts = [
        ('روش', payload.get('generalSolution', '')),
        ('مسئله', payload.get('problem', '')),
        ('راه‌حلِ کتاب', solution),
    ]
    joined = '\n'.join(f'{label}: {text}' for label, text in parts if text).strip()
    return joined or payload.get('text', '')


def _payload_title(payload: dict) -> str:
    return (
        payload.get('lesson_title')
        or payload.get('title')
        or payload.get('type')
        or payload.get('chapter')
        or ''
    )