MathApp/rag/management/commands/rag_model_benchmark.py

"""
مقایسه‌ی مدل‌های چت برای دستیار («از پی بپرس») با عدد، نه با حدس.

چرا جدا از سنجشِ استخراجِ PDF: آنجا مدل صفحه‌ی کتاب را می‌خواند (بینایی)،
اینجا باید مسئله را درست حل کند و مرحله‌به‌مرحله توضیح دهد. مدلی که در یکی
بهترین است لزوماً در دیگری نیست — gpt-4o-mini در استخراج مختصات جعل می‌کرد و
در چت هم مشتقِ x³ − 6x² + 9x را غلط گرفت.

روش:
  ۱. از هر کتاب چند مسئله‌ی بی‌شکل برمی‌دارد که `rag_verify_solutions` جوابِ
     کتاب را تأیید کرده (`agrees`) و جوابِ مستقلش (`checked_answer`) ثبت است؛
     همان جواب کلیدِ تصحیح است.
  ۲. هر مسئله را از همان مسیرِ واقعیِ دستیار می‌گذراند: retrieval روی کلِ کتابِ
     آن پایه (مثلِ صفحه‌ی «از پی بپرس») + همان پرامپتِ تولید. خودِ همان تکه از
     context حذف می‌شود تا مدل جواب را از کتاب رونویسی نکند.
  ۳. یک مدلِ داورِ ثابت فقط جوابِ نهایی را با کلید مقایسه می‌کند.
  ۴. هزینه‌ی واقعی از `/dashboard/billing/usage` قبل و بعد از هر مدل خوانده
     می‌شود — قیمتِ GapGPT با جدولِ رسمی یکی نیست (docs/extraction-playbook.md).

    # رایگان: فقط نمونه را نشان می‌دهد
    python manage.py rag_model_benchmark --dry-run

    python manage.py rag_model_benchmark --models gpt-4o-mini gpt-4.1-mini o4-mini \\
        --per-grade 6 --json /tmp/bench.json

⚠️ GapGPT فقط از داخلِ ایران جواب می‌دهد؛ این را روی سرور اجرا کنید.
"""
import hashlib
import json
import re
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

from django.core.management.base import BaseCommand, CommandError

from rag.management.commands.rag_verify_solutions import (
    AccountLimitReached, is_account_limit, needs_figure,
)
from rag.models import LessonChunk
from rag.services import clients
from rag.services.config import rag_is_configured
from rag.services.generation import completion_kwargs, generate_answer, grade_tier
from rag.services.retrieval import retrieve_context

DEFAULT_MODELS = ['gpt-4o-mini', 'gpt-4.1-mini', 'gpt-4.1', 'o4-mini']
DEFAULT_JUDGE = 'gpt-4.1'

JUDGE_PROMPT = """یک مسئله‌ی ریاضیِ مدرسه، جوابِ درستِ آن (کلید)، و پاسخِ یک معلم به دانش‌آموز
را می‌بینی. فقط بگو جوابِ نهاییِ معلم با کلید یکی است یا نه.

قواعد:
- شکلِ نوشتن مهم نیست: 0.5 = ۱/۲ = نیم؛ «۳/۵» در کلید ممکن است ممیزِ اعشار باشد (۳٫۵).
- ترتیب، واحدِ هم‌ارز، و گِرد کردنِ معقول خطا نیست.
- اگر کلید چند بخش دارد، همه‌ی بخش‌ها باید درست باشند.
- اگر معلم جوابِ نهایی نداده یا چند جوابِ متناقض داده، "wrong".
- اگر خودِ کلید با مسئله نمی‌خواند یا مبهم است، "unclear".

مسئله:
{problem}

کلید:
{key}

پاسخِ معلم:
{answer}

فقط JSON: {{"verdict": "correct|wrong|unclear", "why": "حداکثر ۸۰ نویسه"}}"""

# نشانه‌ی LaTeX/markdown در خروجیِ *خامِ* مدل (پیش از تورِ ایمنی)
_STEP_LINE = re.compile(r'^\s*([0-9۰-۹]+[)\.)]|مرحله)', re.M)


_KEY_NOT_AN_ANSWER = re.compile(r'کتاب|اثبات|درست است|صحیح است|مطابق|توضیح')


def _gradable(chunk) -> bool:
    """
    فقط مسئله‌ای که جوابِ مشخص و کوتاه دارد (عدد، عبارت، مجموعه). کلیدهایی
    مثلِ «راه‌حلِ کتاب مطابقِ تعریف است» یا اثبات‌ها را نمی‌شود با کلید سنجید.
    """
    key = (chunk.checked_answer or '').strip()
    return (
        0 < len(key) <= 80
        and bool(re.search(r'[0-9۰-۹]', key))
        and not _KEY_NOT_AN_ANSWER.search(key)
        and not chunk.problem.strip().startswith(('اثبات', 'ثابت کنید', 'نشان دهید'))
    )


def _stable_order(chunk_id: int, seed: str) -> str:
    return hashlib.sha1(f'{seed}:{chunk_id}'.encode()).hexdigest()


def _grade_of(item) -> str:
    return (item['grade'] or item['collection'].split()[-1]).strip()


def _usage_cents() -> float | None:
    """مصرفِ تجمعیِ حساب به سنت، یا None اگر provider این endpoint را ندارد."""
    import httpx
    from django.conf import settings

    try:
        response = httpx.get(
            f'{settings.OPENAI_BASE_URL.rstrip("/")}/dashboard/billing/usage',
            headers={'Authorization': f'Bearer {settings.OPENAI_API_KEY}'}, timeout=20,
        )
        data = response.json()
        return float(data['total_usage']) if 'total_usage' in data else None
    except Exception:  # noqa: BLE001
        return None


class Command(BaseCommand):
    help = 'مقایسه‌ی دقت/هزینه/سرعتِ مدل‌های چت روی مسئله‌های تأییدشده‌ی کتاب'

    def add_arguments(self, parser):
        parser.add_argument('--models', nargs='+', default=DEFAULT_MODELS)
        parser.add_argument('--judge', default=DEFAULT_JUDGE)
        parser.add_argument('--per-grade', type=int, default=6)
        parser.add_argument('--collections', nargs='*', default=None,
                            help='پیش‌فرض: همه‌ی کتاب‌های ریاضی/حسابان')
        parser.add_argument('--seed', default='bench-v1')
        parser.add_argument('--workers', type=int, default=4)
        parser.add_argument('--no-context', action='store_true',
                            help='بدونِ retrieval — فقط توانِ ریاضیِ خودِ مدل')
        parser.add_argument('--json', dest='json_path', default=None)
        parser.add_argument('--dry-run', action='store_true')
        parser.add_argument('--resume', action='store_true',
                            help='مدل‌هایی که در فایلِ --json تمام شده‌اند دوباره اجرا نشوند')
        # LessonChunk (و کلیدِ checked_answer) فقط در دیتابیسِ محلی است، ولی GapGPT
        # فقط از ایران جواب می‌دهد: نمونه را اینجا بنویس، روی سرور بخوان.
        parser.add_argument('--export-sample', default=None, metavar='PATH')
        parser.add_argument('--sample-file', default=None, metavar='PATH')

    # ------------------------------------------------------------------ sample
    def _sample(self, per_grade, collections, seed):
        qs = LessonChunk.objects.filter(solution_verdict='agrees').exclude(
            checked_answer='').exclude(problem='')
        if collections:
            qs = qs.filter(collection__in=collections)
        else:
            qs = qs.filter(collection__regex=r'^(ریاضی|حسابان) ')
        by_collection = defaultdict(list)
        for chunk in qs.only('id', 'problem', 'checked_answer', 'collection', 'grade_title',
                             'subject', 'chapter'):
            if needs_figure(chunk.problem) or not _gradable(chunk):
                continue
            by_collection[chunk.collection].append(chunk)
        items = []
        for collection in sorted(by_collection):
            chunks = sorted(by_collection[collection], key=lambda c: _stable_order(c.id, seed))
            items.extend({
                'id': c.id, 'problem': c.problem, 'key': c.checked_answer,
                'collection': c.collection, 'grade': c.grade_title,
            } for c in chunks[:per_grade])
        return items

    # ----------------------------------------------------------------- context
    def _context(self, item):
        subject = item['collection'].rsplit(' ', 1)[0]
        hits, _meta = retrieve_context(item['problem'], subject=subject, grade_title=_grade_of(item))
        needle = ' '.join(item['problem'].split())[:60]
        # همان تکه (و هر تکه‌ای که عینِ این مسئله را دارد) بیرون: وگرنه مدل
        # فقط راه‌حلِ کتاب را رونویسی می‌کند و سنجه بی‌معنا می‌شود.
        return [h for h in hits if needle not in ' '.join(h.get('summary', '').split())]

    # ------------------------------------------------------------------ judge
    def _judge(self, judge_model, item, answer):
        client = clients.get_openai_client()
        prompt = JUDGE_PROMPT.format(problem=item['problem'], key=item['key'],
                                     answer=answer[-2500:])
        for attempt in range(3):
            try:
                completion = client.chat.completions.create(
                    model=judge_model, messages=[{'role': 'user', 'content': prompt}],
                    response_format={'type': 'json_object'}, **completion_kwargs(judge_model),
                )
                data = json.loads(completion.choices[0].message.content or '{}')
                verdict = data.get('verdict')
                if verdict in ('correct', 'wrong', 'unclear'):
                    return verdict, str(data.get('why') or '')[:120]
            except Exception as exc:  # noqa: BLE001
                if is_account_limit(exc):
                    raise AccountLimitReached(str(exc)[:160]) from exc
                time.sleep(2 * (attempt + 1))
        return 'unclear', 'داوری شکست خورد'

    # ------------------------------------------------------------------ handle
    def handle(self, *args, **opts):
        if opts['sample_file']:
            with open(opts['sample_file'], encoding='utf-8') as fh:
                items = json.load(fh)
        else:
            items = self._sample(opts['per_grade'], opts['collections'], opts['seed'])
        if opts['export_sample']:
            with open(opts['export_sample'], 'w', encoding='utf-8') as fh:
                json.dump(items, fh, ensure_ascii=False, indent=1)
            self.stdout.write(f'{len(items)} مسئله در {opts["export_sample"]} نوشته شد.')
            return
        if not items:
            raise CommandError('هیچ مسئله‌ی تأییدشده‌ای پیدا نشد.')

        per_collection = defaultdict(int)
        for item in items:
            per_collection[item['collection']] += 1
        self.stdout.write(f'{len(items)} مسئله از {len(per_collection)} کتاب:')
        for collection, n in sorted(per_collection.items()):
            self.stdout.write(f'  {collection:18s} {n}')
        calls = len(items) * len(opts['models'])
        self.stdout.write(f'{calls} پاسخ + {calls} داوری با {opts["judge"]}')
        if opts['dry_run']:
            for item in items[:3]:
                self.stdout.write(f'\n— {item["collection"]}: {item["problem"][:160]}\n  کلید: {item["key"][:80]}')
            return
        if not rag_is_configured():
            raise CommandError('RAG پیکربندی نشده (OPENAI_API_KEY/QDRANT_API_KEY).')

        # context یک بار برای همه‌ی مدل‌ها: مقایسه منصفانه و embedding ارزان‌تر.
        contexts = {}
        if not opts['no_context']:
            self.stdout.write('retrieval …')
            try:
                with ThreadPoolExecutor(opts['workers']) as pool:
                    for item, ctx in zip(items, pool.map(self._context, items)):
                        contexts[item['id']] = ctx
            except Exception as exc:  # noqa: BLE001
                if is_account_limit(exc):
                    raise CommandError(f'⛔ اعتبارِ حساب تمام شده است: {str(exc)[:160]}') from exc
                raise

        # نتیجه‌ی هر مدل همان لحظه روی دیسک می‌رود. اجرای اول وسطِ کار با تمام
        # شدنِ اعتبار مرد و چون فقط آخرِ کار ذخیره می‌کرد، پاسخ‌های چهار مدلِ
        # پرداخت‌شده از دست رفت. با --resume مدل‌های تمام‌شده دوباره اجرا نمی‌شوند.
        results, summary = [], {}
        if opts['resume'] and opts['json_path']:
            try:
                with open(opts['json_path'], encoding='utf-8') as fh:
                    saved = json.load(fh)
                summary = saved.get('summary', {})
                results = [r for r in saved.get('results', []) if r['model'] in summary]
                self.stdout.write(f'ادامه: {len(summary)} مدل از قبل — {", ".join(summary)}')
            except FileNotFoundError:
                pass

        chunks_by_id = {item['id']: item for item in items}
        try:
            for model in opts['models']:
                if model in summary:
                    continue
                self._run_model(model, items, contexts, chunks_by_id, opts, results, summary)
                self._save(opts['json_path'], summary, results)
        except AccountLimitReached as exc:
            self._save(opts['json_path'], summary, results)
            self.stderr.write(self.style.ERROR(
                f'\n⛔ اعتبارِ حساب تمام شد ({exc}). نتیجه‌ی مدل‌های تمام‌شده ذخیره شد؛ '
                'بعد از شارژ با همان دستور و --resume ادامه بدهید.'))
        self._report([m for m in opts['models'] if m in summary], results, summary)

    def _save(self, path, summary, results):
        if not path:
            return
        with open(path, 'w', encoding='utf-8') as fh:
            json.dump({'summary': summary, 'results': results}, fh, ensure_ascii=False, indent=1)

    def _run_model(self, model, items, contexts, chunks_by_id, opts, results, summary):
        self.stdout.write(f'\n▶ {model}')
        before = _usage_cents()
        started = time.monotonic()

        def run(item):
            t0 = time.monotonic()
            try:
                answer, _f, meta = generate_answer(
                    item['problem'], contexts.get(item['id']) or None, _grade_of(item), model=model,
                )
                return item, answer, meta, time.monotonic() - t0, None
            except Exception as exc:  # noqa: BLE001
                if is_account_limit(exc):
                    raise AccountLimitReached(str(exc)[:160]) from exc
                return item, '', {}, time.monotonic() - t0, str(exc)[:200]

        with ThreadPoolExecutor(opts['workers']) as pool:
            answered = list(pool.map(run, items))
        after = _usage_cents()
        wall = time.monotonic() - started
        cost = (after - before) if before is not None and after is not None else None

        rows = []
        for item, answer, meta, seconds, error in answered:
            rows.append({
                'model': model, 'chunkId': item['id'], 'collection': item['collection'],
                'tier': grade_tier(_grade_of(item)), 'problem': item['problem'],
                'key': item['key'], 'answer': answer, 'seconds': round(seconds, 1),
                'sanitized': bool(meta.get('sanitized')),
                'completionTokens': meta.get('completionTokens', 0),
                'hasFinal': '✅' in answer, 'steps': len(_STEP_LINE.findall(answer)),
                'error': error, 'verdict': 'error', 'why': '',
            })
        errors = sum(1 for r in rows if r['error'])
        self.stdout.write(
            f'  {len(rows) - errors} پاسخ، {errors} خطا، '
            f'{"هزینه: %.2f سنت" % cost if cost is not None else "هزینه: نامعلوم"}، {wall:.0f}s'
        )

        todo = [r for r in rows if r['answer']]
        with ThreadPoolExecutor(opts['workers']) as pool:
            verdicts = list(pool.map(
                lambda r: self._judge(opts['judge'], chunks_by_id[r['chunkId']], r['answer']), todo))
        for r, (verdict, why) in zip(todo, verdicts):
            r['verdict'], r['why'] = verdict, why
        results.extend(rows)
        summary[model] = {'costCents': cost, 'wallSeconds': round(wall, 1)}

    def _report(self, models, results, summary):
        tiers = ['primary', 'middle', 'upper']
        header = f'{"مدل":18s} {"درست":>7s} {"دبستان":>7s} {"متوسطه۱":>7s} {"متوسطه۲":>7s} ' \
                 f'{"نامعلوم":>6s} {"سنت/پاسخ":>8s} {"ثانیه":>6s} {"LaTeX":>6s} {"✅":>5s}'
        self.stdout.write('\n' + header)
        for model in models:
            rows = [r for r in results if r['model'] == model]
            judged = [r for r in rows if r['verdict'] in ('correct', 'wrong')]

            def pct(subset):
                graded = [r for r in subset if r['verdict'] in ('correct', 'wrong')]
                if not graded:
                    return '—'
                return f'{100 * sum(r["verdict"] == "correct" for r in graded) / len(graded):.0f}%'

            cost = summary[model]['costCents']
            per = f'{cost / len(rows):.3f}' if cost is not None and rows else '?'
            secs = sum(r['seconds'] for r in rows) / max(len(rows), 1)
            latex = sum(r['sanitized'] for r in rows) / max(len(rows), 1)
            final = sum(r['hasFinal'] for r in rows) / max(len(rows), 1)
            self.stdout.write(
                f'{model:18s} {pct(judged):>7s} '
                + ' '.join(f'{pct([r for r in rows if r["tier"] == t]):>7s}' for t in tiers)
                + f' {sum(r["verdict"] == "unclear" for r in rows):>6d} {per:>8s} {secs:>6.1f}'
                  f' {100 * latex:>5.0f}% {100 * final:>4.0f}%'
            )