MathApp/rag/management/commands/rag_ingest_book.py

"""
ingestِ یک کتابِ درسیِ کامل از data/textbooks/ به Qdrant.

نمونه:
    python manage.py rag_ingest_book --subject ریاضی --grade نهم --plan
    python manage.py rag_ingest_book --subject ریاضی --grade نهم --chapters 1 --max-pages 3
    python manage.py rag_ingest_book --subject ریاضی --grade نهم

هر اجرا از جایی که مانده ادامه می‌دهد (صفحه‌های ثبت‌شده رد می‌شوند)، پس
قطع کردنش با Ctrl-C بی‌خطر است و اجرای دوباره کارِ رفته را دوباره نمی‌خرد.
"""
from django.core.management.base import BaseCommand, CommandError

from rag.services.clients import RagNotConfigured
from rag.services.config import build_collection_name, rag_is_configured
from rag.services.ingest import _done_pages, estimate_page_cost, ingest_book
from rag.services.textbooks import TEXTBOOKS, get_textbook
from rag.models import IngestedPage, LessonChunk


class Command(BaseCommand):
    help = 'ingestِ کتابِ درسی (PDF) به LessonChunk + Qdrant، فصل‌به‌فصل و قابلِ ازسرگیری.'

    def add_arguments(self, parser):
        parser.add_argument('--subject', default='ریاضی')
        parser.add_argument('--grade', required=True, help='مثلاً: نهم')
        parser.add_argument('--chapters', nargs='*', type=int, default=None,
                            help='شماره‌ی فصل‌ها؛ خالی یعنی همه.')
        parser.add_argument('--max-pages', type=int, default=None,
                            help='سقفِ صفحه در این اجرا (کنترلِ هزینه).')
        parser.add_argument('--dpi', type=int, default=200)
        parser.add_argument('--no-resume', action='store_true',
                            help='صفحه‌های قبلاً ثبت‌شده را هم دوباره پردازش کن (نقطه‌ی تکراری می‌سازد).')
        parser.add_argument('--plan', action='store_true',
                            help='فقط نشان بده چه چیزی پردازش می‌شود؛ هیچ فراخوانیِ پولی نکن.')

    def handle(self, *args, **opts):
        subject, grade = opts['subject'], opts['grade']
        book = get_textbook(subject, grade)
        if book is None:
            known = '، '.join(sorted(TEXTBOOKS)) or '—'
            raise CommandError(f'کتابی برای «{subject} {grade}» ثبت نشده. موجود: {known}')
        if not book.filename:
            raise CommandError(
                f'«{book.key}» فایلِ PDF ندارد؛ دادهٔ واقعی‌اش از قبل روی Qdrant است.\n'
                f'برای واردکردنش:  python manage.py rag_import_collection '
                f'--collection "{book.key}"'
            )

        collection = build_collection_name(book.subject, book.grade)
        wanted = list(book.chapters) if not opts['chapters'] else [
            c for n in opts['chapters'] if (c := book.chapter(n)) is not None
        ]
        if not wanted:
            raise CommandError('هیچ فصلی با شماره‌های داده‌شده پیدا نشد.')

        if opts['plan']:
            return self._plan(book, collection, wanted)

        if not rag_is_configured():
            raise CommandError(
                'RAG پیکربندی نشده (OPENAI_API_KEY / QDRANT_API_KEY خالی است یا RAG_FORCE_MOCK روشن است).'
            )

        self.stdout.write(self.style.MIGRATE_HEADING(
            f'ingest «{book.key}» → collection «{collection}»'
        ))
        try:
            result = ingest_book(
                subject=book.subject, grade=book.grade,
                chapters=[c.number for c in wanted],
                dpi=opts['dpi'], max_pages=opts['max_pages'],
                resume=not opts['no_resume'], progress=self._progress,
            )
        except (RagNotConfigured, FileNotFoundError, ValueError) as exc:
            raise CommandError(str(exc)) from exc

        self.stdout.write('')
        style = self.style.SUCCESS if result['status'] == 'success' else self.style.WARNING
        self.stdout.write(style(
            f"وضعیت: {result['status']} — {result['pages_processed']} صفحه پردازش شد، "
            f"{result['pages_skipped']} صفحه از قبل موجود بود، "
            f"{result['chunks_created']} chunk ساخته شد."
        ))
        self.stdout.write(
            f"مصرف: {result['prompt_tokens']:,} توکن ورودی + "
            f"{result['completion_tokens']:,} خروجی  →  ≈ ${result['cost_usd']:.2f} "
            f"({result['model']})"
        )
        if result['failed_pages']:
            self.stdout.write(self.style.ERROR(
                'صفحه‌های خطادار (اجرای دوباره سراغشان می‌رود): '
                + '، '.join(str(p) for p in result['failed_pages'])
            ))

    # ------------------------------------------------------------------
    def _plan(self, book, collection, wanted):
        self.stdout.write(self.style.MIGRATE_HEADING(
            f'نقشه‌ی «{book.key}» → collection «{collection}»  (فایل: {book.filename})'
        ))
        total_pages = total_left = 0
        for chapter in wanted:
            first, last = book.pdf_pages(chapter)
            pages = last - first + 1
            done = len(_done_pages(collection, chapter.title))
            left = max(pages - done, 0)
            total_pages += pages
            total_left += left
            self.stdout.write(
                f'  فصل {chapter.number}: {chapter.title}\n'
                f'      صفحه‌ی کتاب {chapter.first_page}–{chapter.last_page}'
                f'  →  صفحه‌ی PDF {first}–{last}   ({pages} صفحه، {done} انجام‌شده، {left} مانده)'
            )
        self.stdout.write('')
        self.stdout.write(f'جمع: {total_pages} صفحه، {total_left} صفحه مانده.')
        from django.conf import settings

        model = settings.OPENAI_VISION_MODEL
        per_page = estimate_page_cost(model)
        self.stdout.write(self.style.WARNING(
            f'هر صفحه یک فراخوانیِ vision دارد ({model}، ~${per_page:.4f} در صفحه، '
            f'اندازه‌گیری‌شده روی ۱۱۱ صفحه‌ی واقعی): '
            f'حدود ${total_left * per_page:.2f} برای صفحه‌های مانده '
            f'(با Batch API نصف: ${total_left * per_page / 2:.2f}).'
        ))

    def _progress(self, event):
        kind = event.get('kind')
        if kind == 'chapter':
            first, last = event['pages']
            self.stdout.write(self.style.HTTP_INFO(
                f"\n▸ فصل {event['chapter']}: {event['title']}  "
                f"(PDF {first}–{last}، {event['already']} صفحه از قبل موجود)"
            ))
        elif kind == 'page':
            self.stdout.write(f"    صفحه {event['page']}: {event['chunks']} chunk")
        elif kind == 'page_failed':
            self.stdout.write(self.style.ERROR(
                f"    صفحه {event['page']}: خطا — {event['error'][:120]}"
            ))
        elif kind == 'budget_reached':
            self.stdout.write(self.style.WARNING('    سقفِ --max-pages پر شد؛ توقف.'))