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

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
        emb_literal = "[" + ",".join(str(x) for x in emb) + "]"
        rows = db.execute(sa_text(f"""
            SELECT id, 1 - (embedding <=> '{emb_literal}'::vector) AS sim
            FROM questions
            WHERE embedding IS NOT NULL AND id != {question.id}
            ORDER BY embedding <=> '{emb_literal}'::vector
            LIMIT {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 helping a student understand the following question. "
        "Be accurate, educational, and concise. You may reveal and explain the correct answer.\n\n"
        f"=== Current 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 questions about this topic. "
        "If they ask why an option is wrong, explain the underlying concept. "
        "Keep responses focused and under 200 words unless a longer explanation is needed."
    )
    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."""
    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)}")