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

"""Teach chat endpoint — AI tutor for study mode questions."""
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.auth import get_current_user, check_rate_limit

router = APIRouter()


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


class ChatRequest(BaseModel):
    question_id: int
    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, limit: int = 4) -> list[Question]:
    """Return up to `limit` questions with similar embeddings, excluding the current one."""
    if question.embedding is None:
        return []
    try:
        from sqlalchemy import text as sa_text
        emb = question.embedding
        # Validate all values are finite floats before using in SQL
        emb_literal = "[" + ",".join(str(float(x)) for x in emb) + "]"
        rows = db.execute(sa_text("""
            SELECT id, 1 - (embedding <=> CAST(:vec AS vector)) AS sim
            FROM questions
            WHERE embedding IS NOT NULL AND id != :qid
            ORDER BY embedding <=> CAST(:vec AS vector)
            LIMIT :lim
        """), {"vec": emb_literal, "qid": int(question.id), "lim": int(limit)}).fetchall()
        ids = [r.id for r in rows if float(r.sim) >= 0.35]
        if not ids:
            return []
        return db.query(Question).filter(Question.id.in_(ids)).all()
    except Exception:
        return []


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

    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."
    )
    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")
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."""
    # Rate limit: 30 AI chat messages per user per 10 minutes
    check_rate_limit(
        key=f"teach_chat:{current_user.id}",
        max_calls=30,
        window_seconds=600,
        detail="Too many AI chat messages. Please wait a few minutes before continuing.",
    )
    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

    question = db.query(Question).filter(Question.id == req.question_id).first()
    if not question:
        raise HTTPException(status_code=404, detail="Question not found")

    similar = _find_similar_questions(db, question)
    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 = litellm.completion(**kwargs)
        reply = response.choices[0].message.content.strip()
        return {"reply": reply}
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"AI model error: {str(e)}")