202 lines
7.9 KiB
Python
202 lines
7.9 KiB
Python
"""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):
|
|
model_config = {"protected_namespaces": ()}
|
|
|
|
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.\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."""
|
|
# Rate limit: 30 AI chat messages per user per 10 minutes (admins/unthrottled users exempt)
|
|
check_rate_limit(
|
|
key=f"teach_chat:{current_user.id}",
|
|
max_calls=30,
|
|
window_seconds=600,
|
|
detail="You've sent too many messages to the AI tutor. Please wait a few minutes before continuing. If you need this limit raised, contact an admin.",
|
|
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
|
|
|
|
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 = 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.")
|