MathApp/rag/services/generation.py

"""
مرحله‌ی generation: پرسش + context بازیابی‌شده → پاسخِ فارسی + چند پیشنهادِ
ادامه.

اگه context خالی باشه، مدل به‌عنوانِ معلمِ عمومی جواب می‌ده (بدونِ ادعای
استناد به کتاب) — این حالت وقتی پیش می‌آد که subject/lessonNumber ارسال
نشده یا اون درس هنوز در Qdrant نیست.
"""
import logging
import time
from collections.abc import Iterator

from django.conf import settings

from . import clients

logger = logging.getLogger(__name__)

_DEFAULT_FOLLOW_UPS = ['یک مثال دیگه بزن', 'بیشتر توضیح بده']

# ⚠️ با هر تغییرِ معنادار در متنِ پرامپت‌های پایین، این را یکی جلو ببرید.
# روی هر پیامِ دستیار ذخیره می‌شود تا وقتی لایک/دیسلایکِ کاربرها جمع شد،
# بشود گفت «کدام نسخه‌ی پرامپت جوابِ بهتری می‌داد» — بدونِ این، همه‌ی فیدبک‌ها
# در یک کاسه می‌ریزند و مقایسه‌ی نسخه‌ها ممکن نیست.
PROMPT_VERSION = 'tutor-v2-grade-aware'

# پایه‌ها تا ششم دبستان‌اند، هفتم تا نهم متوسطه‌ی اول و دهم به بعد متوسطه‌ی دوم.
_SECONDARY_GRADES = {'هفتم', 'هشتم', 'نهم', 'دهم', 'یازدهم', 'دوازدهم'}


def _persona(grade_title: str) -> tuple[str, str]:
    """(توصیفِ معلم، توصیفِ دانش‌آموز) متناسب با پایه."""
    grade = (grade_title or '').strip()
    if grade in _SECONDARY_GRADES:
        return f'معلمِ ریاضیِ پایه‌ی {grade}', f'دانش‌آموزِ پایه‌ی {grade}'
    if grade:
        return f'معلمِ پایه‌ی {grade} دبستان', f'دانش‌آموزِ پایه‌ی {grade}'
    return 'معلمِ فارسی‌زبان', 'دانش‌آموز'


def system_prompt(grade_title: str = '') -> str:
    """
    قبلاً «معلمِ دبستانی» ثابت بود. برای پایه‌های نهم/دهم این هم لحن را بچگانه
    می‌کرد و هم مدل را وادار می‌کرد مفاهیمی مثل اتحاد، مثلثات و تابع را تا حدِ
    دبستان ساده کند — یعنی دقیقاً همان‌جایی که بچه‌ها گیر می‌کنند بی‌فایده
    می‌شد. حالا لحن از پایه می‌آید.
    """
    teacher, _ = _persona(grade_title)
    return (
        f'تو یک {teacher} و فارسی‌زبان هستی. پاسخ‌ها را روشن، دقیق و '
        'گام‌به‌گام می‌دهی و از نمادها و اصطلاح‌های ریاضیِ همان پایه استفاده '
        'می‌کنی؛ نه ساده‌تر از سطحِ کتاب، نه پیچیده‌تر از آن.'
    )


def _prompt_with_context(question: str, contexts: list[str], grade_title: str = '') -> str:
    joined = '\n\n---\n\n'.join(contexts)
    teacher, student = _persona(grade_title)
    return f"""
    تو {teacher} هستی. با تکیه بر «متنِ درس» زیر و با زبانِ خودت، به سوالِ
    {student} پاسخِ روشن و گام‌به‌گام بده. لازم نیست پاسخ عیناً در متن باشد؛
    می‌توانی از اطلاعاتِ متن نتیجه‌گیری و توضیح بدهی.
    اگر مسئله چند مرحله دارد، مرحله‌ها را جدا و شماره‌گذاری‌شده بنویس.
    فقط اگر متنِ درس آشکارا درباره‌ی موضوعِ دیگری‌ست و هیچ ربطی به سوال ندارد،
    بنویس: «این موضوع در این درس پوشش داده نشده است.»

    متنِ درس:
    {joined}

    سوالِ {student}: {question}
    """


def _prompt_without_context(question: str, grade_title: str = '') -> str:
    teacher, student = _persona(grade_title)
    return f"""
    تو {teacher} هستی. به این سوالِ یک {student} به فارسی، روشن و گام‌به‌گام
    پاسخ بده.

    سوال: {question}
    """


def generate_answer(question: str, context_chunks: list[dict] | None,
                    grade_title: str = '') -> tuple[str, list[str], dict]:
    """
    context_chunks: خروجیِ retrieve_context ([{title, summary, score}, ...]) یا None/[].
    خروجی: (answer, follow_ups, meta)

    `meta` همان چیزی‌ست که روی پیامِ دستیار ذخیره می‌شود: مدل، نسخه‌ی پرامپت،
    توکنِ مصرف‌شده و مدتِ تولید. بدونِ اینها نه می‌شود فهمید کدام نسخه‌ی پرامپت
    بهتر جواب می‌دهد، نه اینکه چقدر از سهمیه‌ی API خرج شده است.
    """
    contexts = [c['summary'] for c in (context_chunks or []) if c.get('summary')]
    prompt = (
        _prompt_with_context(question, contexts, grade_title) if contexts
        else _prompt_without_context(question, grade_title)
    )

    client = clients.get_openai_client()
    started = time.monotonic()
    completion = client.chat.completions.create(
        model=settings.OPENAI_CHAT_MODEL,
        messages=[
            {'role': 'system', 'content': system_prompt(grade_title)},
            {'role': 'user', 'content': prompt},
        ],
        temperature=0.3,
    )
    latency_ms = int((time.monotonic() - started) * 1000)

    answer = (completion.choices[0].message.content or '').strip()
    if not answer:
        answer = 'متأسفم، الان نتونستم پاسخِ مناسبی بسازم. لطفاً دوباره بپرس.'

    usage = getattr(completion, 'usage', None)
    meta = {
        'model': settings.OPENAI_CHAT_MODEL,
        'promptVersion': PROMPT_VERSION,
        'latencyMs': latency_ms,
        'promptTokens': getattr(usage, 'prompt_tokens', None) or 0,
        'completionTokens': getattr(usage, 'completion_tokens', None) or 0,
        'hadContext': bool(contexts),
    }
    return answer, list(_DEFAULT_FOLLOW_UPS), meta


def stream_answer(question: str, context_chunks: list[dict] | None,
                  grade_title: str = '') -> Iterator[str]:
    """
    همون generate_answer ولی توکن‌به‌توکن (برای SSE). فقط قطعه‌های متنِ پاسخ
    را yield می‌کند؛ follow_upها ثابت‌اند (_DEFAULT_FOLLOW_UPS) و caller خودش
    اضافه‌شان می‌کند.
    """
    contexts = [c['summary'] for c in (context_chunks or []) if c.get('summary')]
    prompt = (
        _prompt_with_context(question, contexts, grade_title) if contexts
        else _prompt_without_context(question, grade_title)
    )

    client = clients.get_openai_client()
    stream = client.chat.completions.create(
        model=settings.OPENAI_CHAT_MODEL,
        messages=[
            {'role': 'system', 'content': system_prompt(grade_title)},
            {'role': 'user', 'content': prompt},
        ],
        temperature=0.3,
        stream=True,
    )
    for event in stream:
        if not event.choices:
            continue
        delta = event.choices[0].delta.content or ''
        if delta:
            yield delta


def default_follow_ups() -> list[str]:
    return list(_DEFAULT_FOLLOW_UPS)