MathApp/rag/services/ingest.py

"""
هماهنگ‌کننده‌ی ingestion: متن/PDF درس → chunk/خلاصه/embedding → ذخیره در
DB (LessonChunk، منبعِ حقیقت) + upsert به Qdrant (index).

نامِ collection و کلیدهای payload عمداً همون قراردادِ Rag_backend را دنبال
می‌کنن تا با دادهٔ موجودِ تیمِ دیگه سازگار بمونه.
"""
import logging
import os
import tempfile

from django.db import transaction

from ..models import LessonChunk
from .chunking import chunk_and_summarize
from .config import build_collection_name, resolve_grade_title
from .embedding import embed_text
from .known_collections import is_math_subject, resolve_math_chapter
from . import pdf as pdf_service
from .qdrant_store import ensure_collection, upsert_points

logger = logging.getLogger(__name__)


def _payload(chunk: LessonChunk) -> dict:
    """
    payloadِ نقطه‌ی Qdrant.

    ⚠️ کلیدِ `chapter` (رشته‌ی کاملِ فصل) حتماً باید نوشته بشه: جستجوی ریاضی در
    retrieval.py دقیقاً روی همین کلید فیلتر می‌کنه (نه `chapter_number`)، پس
    بدونش هر چیزی که ingest کنیم برای همیشه نامرئی می‌مونه — بدونِ هیچ خطایی.
    بقیه‌ی کلیدهای ساختاری هم عمداً هم‌نامِ قراردادِ تیمِ دیگه‌ان تا collectionِ
    ما و اونا یک شکل بمونن.
    """
    payload = {
        'lesson_id': chunk.lesson_id,
        'lesson_number': chunk.lesson_number,
        'chapter_number': chunk.chapter,
        'chunk_order': chunk.order,
        'lesson_title': chunk.title,
        'summary': chunk.summary,
    }
    for key, value in (
        ('chapter', chunk.chapter_title),
        ('grade', chunk.grade_title),
        ('type', chunk.problem_type),
        ('keywords', chunk.keywords),
        ('generalSolution', chunk.general_solution),
        ('problem', chunk.problem),
        ('solution', chunk.solution),
    ):
        if value:
            payload[key] = value
    return payload


def _persist_and_upsert(rows: list[dict], *, lesson_id, grade_id, subject,
                        lesson_number, chapter, title, collection) -> list[LessonChunk]:
    ensure_collection(collection)
    created: list[LessonChunk] = []
    points: list[dict] = []

    with transaction.atomic():
        for row in rows:
            vector = embed_text(row['summary'])
            chunk = LessonChunk.objects.create(
                lesson_id=lesson_id,
                grade_id=grade_id,
                subject=subject,
                lesson_number=lesson_number,
                chapter=chapter,
                order=row['order'],
                title=title,
                original_text=row['original_text'],
                summary=row['summary'],
                embedding_vector=vector,
                collection=collection,
                # فیلدهای ساختاری فقط از مسیرِ PDF/vision می‌آن؛ ingest_text
                # اون‌ها رو نداره و خالی می‌مونن (payload هم اون‌وقت نمی‌نویسدشون).
                **row.get('structured', {}),
            )
            created.append(chunk)
            points.append({'id': chunk.qdrant_point_id, 'vector': vector, 'payload': _payload(chunk)})

    upsert_points(collection, points)
    return created


def ingest_text(*, lesson_id: int, grade_id, subject: str, lesson_number: int,
                text: str, chapter: int | None = None, title: str = '') -> dict:
    grade_title = resolve_grade_title(grade_id)
    collection = build_collection_name(subject, grade_title)
    title = title or f'lesson_{lesson_number}'

    rows = chunk_and_summarize(text)
    if not rows:
        return {'status': 'empty', 'chunks_created': 0, 'collection': collection}

    # برای ریاضی، فیلترِ جستجو رشته‌ی کاملِ فصله؛ از شماره‌ی فصل درش می‌آریم تا
    # محتوایی که این‌طوری ingest می‌شه هم قابلِ پیدا شدن باشه. اگه فصل ناشناخته
    # باشه (پایه‌ای که هنوز در known_collections نیست) خالی می‌مونه.
    structured = {'grade_title': grade_title}
    if is_math_subject(subject) and chapter is not None:
        structured['chapter_title'] = resolve_math_chapter(subject, grade_title, chapter) or ''
    for row in rows:
        row.setdefault('structured', structured)

    created = _persist_and_upsert(
        rows, lesson_id=lesson_id, grade_id=grade_id, subject=subject,
        lesson_number=lesson_number, chapter=chapter, title=title, collection=collection,
    )
    return {
        'status': 'success',
        'collection': collection,
        'chunks_created': len(created),
        'chunk_ids': [c.id for c in created],
    }


def ingest_pdf(*, lesson_id: int, grade_id, subject: str, lesson_number: int,
               pdf_bytes: bytes, chapter: int | None = None, title: str = '',
               chapter_name: str = '') -> dict:
    """
    هر صفحه‌ی PDF → عکس → استخراجِ انواعِ مسئله با vision → هر problemType یک
    chunk. برای درس‌های متنی به‌جای این از ingest_text استفاده کن.
    """
    grade_title = resolve_grade_title(grade_id)
    collection = build_collection_name(subject, grade_title)
    title = title or f'lesson_{lesson_number}'
    chapter_name = chapter_name or title

    fd, tmp_path = tempfile.mkstemp(suffix='.pdf')
    try:
        with os.fdopen(fd, 'wb') as fh:
            fh.write(pdf_bytes)

        start_page, detected_grade = pdf_service.detect_grade_and_start_page(tmp_path)
        images = pdf_service.extract_page_images(tmp_path, first_page=start_page)

        rows: list[dict] = []
        order = 0
        for image in images:
            try:
                parsed = pdf_service.extract_problem_types(image, detected_grade or grade_title, chapter_name)
            except Exception:
                logger.exception('rag.ingest_pdf: تحلیلِ یک صفحه شکست خورد؛ رد شد')
                continue
            page_chapter = (parsed.get('chapter') or '').strip() or chapter_name
            page_grade = (parsed.get('grade') or '').strip() or (detected_grade or grade_title)
            for problem_type in parsed.get('problemTypes', []) or []:
                text = pdf_service.problem_type_to_text(problem_type)
                if not text:
                    continue
                order += 1
                example = problem_type.get('exampleFromImage', {}) or {}
                rows.append({
                    'order': order,
                    'original_text': text,
                    'summary': text,
                    # ساختار رو نگه می‌داریم، نه فقط متنِ صاف‌شده‌اش را: هم
                    # `chapter` برای فیلترِ جستجو لازمه، هم problem/solution
                    # بعداً ورودیِ ساختِ سؤال‌های چهارگزینه‌ای می‌شن.
                    'structured': {
                        'chapter_title': page_chapter,
                        'grade_title': page_grade,
                        'problem_type': (problem_type.get('type') or '').strip(),
                        'keywords': problem_type.get('keywords') or [],
                        'general_solution': (problem_type.get('generalSolution') or '').strip(),
                        'problem': (example.get('problem') or '').strip(),
                        'solution': (example.get('solution') or '').strip(),
                    },
                })
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass

    if not rows:
        return {'status': 'empty', 'chunks_created': 0, 'collection': collection}

    created = _persist_and_upsert(
        rows, lesson_id=lesson_id, grade_id=grade_id, subject=subject,
        lesson_number=lesson_number, chapter=chapter, title=title, collection=collection,
    )
    return {
        'status': 'success',
        'collection': collection,
        'chunks_created': len(created),
        'chunk_ids': [c.id for c in created],
    }