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 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)
    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`
    """
    if payload.get('summary'):
        return payload['summary']
    parts = [
        payload.get('generalSolution', ''),
        payload.get('problem', ''),
        payload.get('solution', ''),
    ]
    joined = '\n'.join(p for p in parts if p).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 ''
    )