pdf-quiz-generator/backend/app/routers/teach.py

"""Teach chat endpoint — AI tutor for study mode questions."""
from datetime import datetime, time, timedelta, timezone

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy.orm import Session

from app.database import get_db
from app.models.question import Question
from app.models.ai_model_config import AIModelConfig
from app.models.user import User
from app.utils.quiz_access import require_question_access
from app.services.quiz_builder import bank_question_predicate
from app.utils.auth import get_current_user, check_rate_limit

router = APIRouter()


def _daily_teach_limit() -> tuple[str, int]:
    now = datetime.now(timezone.utc)
    reset_at = datetime.combine(now.date() + timedelta(days=1), time.min, tzinfo=timezone.utc)
    ttl = max(60, int((reset_at - now).total_seconds()) + 300)
    return now.date().isoformat(), ttl


class ChatMessage(BaseModel):
    role: str  # "user" | "assistant"
    content: str


class ChatRequest(BaseModel):
    model_config = {"protected_namespaces": ()}

    question_id: int
    attempt_id: int | None = None
    messages: list[ChatMessage]
    model_id: int | None = None  # AIModelConfig.id — if None, use default


def _get_teach_model(db: Session, model_config_id: int | None = None):
    """Return (model_id, api_key) for the requested (or default) teach model, or None."""
    if model_config_id:
        m = db.query(AIModelConfig).filter(
            AIModelConfig.id == model_config_id,
            AIModelConfig.task == "teach",
            AIModelConfig.is_active == True,
        ).first()
        if m:
            return (m.model_id, m.api_key or None)
    # Fall back to default, then any active
    m = db.query(AIModelConfig).filter(
        AIModelConfig.task == "teach",
        AIModelConfig.is_active == True,
        AIModelConfig.is_default == True,
    ).first()
    if not m:
        m = db.query(AIModelConfig).filter(
            AIModelConfig.task == "teach",
            AIModelConfig.is_active == True,
        ).first()
    return (m.model_id, m.api_key or None) if m else None


def _find_similar_questions(db: Session, question: Question, user: User, limit: int = 4) -> list[Question]:
    """Filter eligible context in SQL before similarity ranking and LIMIT."""
    if question.embedding is None:
        return []
    try:
        distance = Question.embedding.cosine_distance(question.embedding)
        return db.query(Question).filter(
            bank_question_predicate(user), Question.id != question.id,
            Question.embedding.isnot(None), distance <= 0.65,
        ).order_by(distance).limit(limit).all()
    except Exception:
        return []


def _build_system_prompt(question: Question, similar: list[Question]) -> str:
    opts = ""
    if question.options:
        opts = "\n".join(f"  {i}) {opt}" for i, opt in enumerate(question.options, 1))

    prompt = (
        "You are a medical education tutor. A student is studying the question below.\n"
        "Rules:\n"
        "- Answer the student's question directly. Do NOT ask clarifying questions.\n"
        "- You may reveal and explain the correct answer and why wrong options are wrong.\n"
        "- Use markdown formatting: bold key terms, bullet lists for comparisons.\n"
        "- Keep responses under 200 words unless a detailed explanation is needed.\n"
        "- Never ask 'what would you like to know?' — just explain.\n\n"
        f"=== Question ===\n{question.question_text}\n"
    )
    if opts:
        prompt += f"Options:\n{opts}\n"
    if question.correct_answer:
        prompt += f"Correct Answer: {question.correct_answer}\n"
    if question.explanation:
        prompt += f"Explanation: {question.explanation}\n"

    if similar:
        prompt += "\n=== Related Questions (for broader context) ===\n"
        for i, sq in enumerate(similar, 1):
            prompt += f"{i}. {sq.question_text}"
            if sq.correct_answer:
                prompt += f" → Answer: {sq.correct_answer}"
            prompt += "\n"

    prompt += (
        "\nAnswer the student's question directly. Explain the correct answer, "
        "the underlying concept, and why wrong options are incorrect if relevant. "
        "Do not ask what they want to know — just teach.\n\n"
        "After your explanation, suggest exactly 3 follow-up questions the student might want to ask next. "
        "Put them at the very end, each on its own line starting with '> '. Example:\n"
        "> Why is dopamine not the first-line treatment here?\n"
        "> What are the other causes of this presentation?\n"
        "> How does this differ in neonates vs older children?"
    )
    return prompt


@router.get("/models")
def list_teach_models(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Return available teach models for the frontend to display."""
    models = db.query(AIModelConfig).filter(
        AIModelConfig.task == "teach",
        AIModelConfig.is_active == True,
    ).all()
    return [{"id": m.id, "name": m.name, "model_id": m.model_id, "is_default": m.is_default} for m in models]


@router.post("/chat")
async def chat(
    req: ChatRequest,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Send a message to the teach AI with full question context."""
    question = db.query(Question).filter(Question.id == req.question_id).first()
    require_question_access(db, question, current_user, req.attempt_id, review=True)

    # Daily AI coach quota. Admins, moderators, and unthrottled users are exempt.
    quota_day, quota_ttl = _daily_teach_limit()
    check_rate_limit(
        key=f"teach_chat_daily:{current_user.id}:{quota_day}",
        max_calls=30,
        window_seconds=quota_ttl,
        detail="You've reached today's AI coach limit of 30 messages. Try again tomorrow, or contact an admin if you need the limit raised.",
        user=current_user,
    )
    model_info = _get_teach_model(db, req.model_id)
    if not model_info:
        raise HTTPException(
            status_code=503,
            detail="No teaching AI model is configured. Ask an admin to add a model with task 'teach'.",
        )
    model_id, api_key = model_info

    similar = _find_similar_questions(db, question, current_user)
    system_prompt = _build_system_prompt(question, similar)

    messages = [{"role": "system", "content": system_prompt}]
    for msg in req.messages:
        if msg.role not in ("user", "assistant"):
            continue
        messages.append({"role": msg.role, "content": msg.content})

    try:
        import litellm
        from app.config import settings
        from app.services.ai_service import _proxy_model
        use_model = _proxy_model(model_id)
        kwargs = {
            "model": use_model,
            "messages": messages,
            "max_tokens": 600,
            "temperature": 0.4,
        }
        if api_key:
            kwargs["api_key"] = api_key
        elif settings.LITELLM_API_KEY:
            kwargs["api_key"] = settings.LITELLM_API_KEY
        if settings.LITELLM_API_BASE:
            kwargs["api_base"] = settings.LITELLM_API_BASE
        response = await litellm.acompletion(**kwargs)
        raw = response.choices[0].message.content.strip()

        # Parse out follow-up suggestions (lines starting with "> ")
        lines = raw.splitlines()
        suggestions = [l[2:].strip() for l in lines if l.startswith("> ")]
        reply_lines = [l for l in lines if not l.startswith("> ")]
        # Trim trailing blank lines from reply
        while reply_lines and not reply_lines[-1].strip():
            reply_lines.pop()
        reply = "\n".join(reply_lines).strip()

        return {"reply": reply, "suggestions": suggestions[:3]}
    except Exception as e:
        import logging
        logging.getLogger(__name__).error(f"TeachChat error for user {current_user.id} model {model_id}: {e}")
        raise HTTPException(status_code=502, detail="The AI tutor is temporarily unavailable. Please try again in a moment.")