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

"""AI Mode — a chat that can only answer from the learner's own library.

The router is thin on purpose. Everything that matters is in
`ai_mode_service`: retrieval builds the shortlist, the shortlist is the whole
prompt, and every citation the model writes is checked against that shortlist
before anyone sees it.
"""
import logging
from datetime import datetime, timezone

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

from app.database import get_db
from app.models.conversation import Conversation, ConversationMessage
from app.models.user import User
from app.services import ai_mode_service
from app.services.ai_service import get_model_for_task
from app.utils.auth import check_rate_limit, get_current_user

router = APIRouter()
log = logging.getLogger(__name__)

DAILY_LIMIT = 60
# How much of the thread goes back to the model. Long enough to follow a
# conversation, short enough that the sources stay the bulk of the prompt.
HISTORY_TURNS = 8


def _own(db: Session, conversation_id: int, user: User) -> Conversation:
    conversation = db.get(Conversation, conversation_id)
    if not conversation or conversation.user_id != user.id:
        # Not "forbidden": whether somebody else's thread exists is not this
        # user's business either.
        raise HTTPException(404, "Conversation not found")
    return conversation


def _message_json(message: ConversationMessage) -> dict:
    return {
        "id": message.id, "role": message.role, "content": message.content,
        "citations": message.citations or [],
        "created_at": message.created_at,
    }


@router.get("/conversations")
def list_conversations(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    rows = db.query(Conversation).filter(Conversation.user_id == current_user.id).order_by(
        Conversation.updated_at.desc(), Conversation.id.desc()).limit(50).all()
    return [{"id": c.id, "title": c.title, "updated_at": c.updated_at,
             "message_count": len(c.messages)} for c in rows]


@router.post("/conversations", status_code=201)
def create_conversation(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    conversation = Conversation(user_id=current_user.id, title="New chat")
    db.add(conversation)
    db.commit()
    return {"id": conversation.id, "title": conversation.title, "messages": []}


@router.get("/conversations/{conversation_id}")
def get_conversation(conversation_id: int, db: Session = Depends(get_db),
                     current_user: User = Depends(get_current_user)):
    conversation = _own(db, conversation_id, current_user)
    return {"id": conversation.id, "title": conversation.title,
            "messages": [_message_json(m) for m in conversation.messages]}


class ConversationRename(BaseModel):
    title: str = Field(min_length=1, max_length=200)


@router.patch("/conversations/{conversation_id}")
def rename_conversation(conversation_id: int, data: ConversationRename,
                        db: Session = Depends(get_db),
                        current_user: User = Depends(get_current_user)):
    conversation = _own(db, conversation_id, current_user)
    conversation.title = data.title.strip()
    db.commit()
    return {"id": conversation.id, "title": conversation.title}


@router.delete("/conversations/{conversation_id}", status_code=204)
def delete_conversation(conversation_id: int, db: Session = Depends(get_db),
                        current_user: User = Depends(get_current_user)):
    db.delete(_own(db, conversation_id, current_user))
    db.commit()


class AskIn(BaseModel):
    message: str = Field(min_length=1, max_length=2000)


@router.post("/conversations/{conversation_id}/messages")
async def ask(conversation_id: int, data: AskIn, db: Session = Depends(get_db),
              current_user: User = Depends(get_current_user)):
    """Answer from the learner's library, citing only what retrieval found."""
    conversation = _own(db, conversation_id, current_user)
    question = data.message.strip()

    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    check_rate_limit(
        key=f"ai_mode_daily:{current_user.id}:{today}",
        max_calls=DAILY_LIMIT, window_seconds=86400,
        detail=f"You've reached today's AI Mode limit of {DAILY_LIMIT} messages.",
        user=current_user,
    )

    model_id, api_key = get_model_for_task(db, "teach")
    if not model_id:
        raise HTTPException(503, "No AI model is configured. Ask an admin to set one up.")

    sources = ai_mode_service.retrieve(db, current_user, question)
    history = [{"role": m.role, "content": m.content}
               for m in conversation.messages[-HISTORY_TURNS:]]
    messages = [{"role": "system", "content": ai_mode_service.build_prompt(sources)},
                *history, {"role": "user", "content": question}]

    try:
        import litellm

        from app.config import settings
        from app.services.ai_service import _proxy_model

        kwargs = {"model": _proxy_model(model_id), "messages": messages,
                  "max_tokens": 700, "temperature": 0.3}
        if api_key or settings.LITELLM_API_KEY:
            kwargs["api_key"] = api_key or 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 or "").strip()
    except Exception:
        log.error("AI Mode failed for user %s", current_user.id, exc_info=True)
        raise HTTPException(502, "AI Mode is temporarily unavailable. Try again in a moment.")

    # The safety step: anything the model cited that retrieval did not find is
    # removed here, before it is stored or shown.
    reply, citations = ai_mode_service.enforce_citations(raw, sources)

    db.add(ConversationMessage(conversation_id=conversation.id, role="user",
                               content=question, citations=[]))
    answer = ConversationMessage(conversation_id=conversation.id, role="assistant",
                                 content=reply, citations=citations)
    db.add(answer)
    # The first question names the thread; "New chat" ages badly in a rail of them.
    if conversation.title == "New chat":
        conversation.title = question[:80] + ("…" if len(question) > 80 else "")
    conversation.updated_at = datetime.utcnow()
    db.commit()
    db.refresh(answer)
    return {"message": _message_json(answer), "title": conversation.title,
            "source_count": len(sources)}