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 . import pdf as pdf_service
from .qdrant_store import ensure_collection, upsert_points

logger = logging.getLogger(__name__)


def _payload(chunk: LessonChunk) -> dict:
    return {
        '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,
    }


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,
            )
            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}

    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
            for problem_type in parsed.get('problemTypes', []) or []:
                text = pdf_service.problem_type_to_text(problem_type)
                if not text:
                    continue
                order += 1
                rows.append({'order': order, 'original_text': text, 'summary': text})
    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],
    }