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.services.quiz_builder import GenerateTestRequest, generate_test
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)


class PracticeIn(BaseModel):
    """Which turn to practise. Absent means the latest answer in the thread."""
    message_id: int | None = None
    # The learner's clock is the one in the name — "1 AM" has to mean the hour
    # they remember — so the client sends it rather than the server stamping it.
    title: str | None = None


@router.post("/conversations/{conversation_id}/practice")
def practice(conversation_id: int, data: PracticeIn, db: Session = Depends(get_db),
             current_user: User = Depends(get_current_user)):
    """Turn a chat turn into a study session.

    The chat is for finding out what you do not know; the point of finding out
    is to go and practise it. This is the step between — it builds a session
    from what the answer actually cited, so the questions follow from the
    conversation rather than from the topic in general.

    Study mode, never exam: this is reading followed by practice, not a paper.
    """
    conversation = _own(db, conversation_id, current_user)
    answer = None
    if data.message_id:
        answer = next((m for m in conversation.messages if m.id == data.message_id), None)
        if answer is None:
            raise HTTPException(404, "That message is not in this conversation")
    else:
        answer = next((m for m in reversed(conversation.messages) if m.role == "assistant"), None)
    if answer is None:
        raise HTTPException(400, "There is nothing to practise in this conversation yet")

    # What was asked, as the learner put it — the turn before the answer.
    asked = ""
    for message in conversation.messages:
        if message.id == answer.id:
            break
        if message.role == "user":
            asked = message.content

    ids = ai_mode_service.practice_ids(db, current_user, answer.citations or [], asked)
    if not ids:
        raise HTTPException(404, "No questions in your bank match this conversation yet")

    # Named like every other session. It used to take the conversation's own
    # title, so asking "hi" produced a session called "hi — practice" sitting
    # in a list of "Custom test from Sep 12, 1 AM". The client sends the name
    # because the hour in it has to be the learner's, not the server's.
    title = (data.title or "").strip()[:180] or f"AI Mode session {datetime.utcnow():%b %-d}"
    quiz = generate_test(db, current_user, GenerateTestRequest(
        title=title, category_ids=[], state="all",
        count=len(ids), mode="learning", explicit_ids=ids))
    return {"quiz_id": quiz["id"], "count": len(ids)}


@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)
    # How close the nearest thing in the library actually is, which decides
    # which of the three answers this question gets. The shortlist alone cannot
    # tell you: reciprocal-rank fusion throws the distances away and returns an
    # order that is never empty, so a question about photosynthesis came back
    # with six paediatric sources and an instruction to answer only from them.
    similarity = ai_mode_service.closeness(db, question)
    mode = ai_mode_service.answer_mode(similarity, sources)
    if mode == "open":
        # Nothing to cite, so nothing is offered for citation — the shortlist is
        # not passed to a model that has just been told the library does not
        # cover this.
        sources = []
    history = [{"role": m.role, "content": m.content}
               for m in conversation.messages[-HISTORY_TURNS:]]
    messages = [{"role": "system", "content": ai_mode_service.build_prompt(sources, mode)},
                *history, {"role": "user", "content": question}]

    try:
        from app.services.ai_service import achat

        raw = (await achat(
            model=model_id, messages=messages, max_tokens=700, temperature=0.3,
            api_key=api_key) 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)}