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 hashlib
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.")

    # A greeting is not a query, and searching a clinical corpus for one comes
    # back full of confident nonsense. Decided before retrieval so nothing is
    # searched for, rather than searched for and then thrown away.
    if ai_mode_service.is_small_talk(question):
        sources, similarity, mode = [], None, "chat"
    else:
        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, question)
    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}]

    # The same question, the same answer.
    #
    # Retrieval was already deterministic — measured over fifteen runs, the
    # shortlist and its scores were identical every time. Generation was the
    # whole variance: at temperature 0.3 the same sources and the same prompt
    # produced answers differing by 15-70% of their text, and one differential
    # swung between a 35-word uncited paraphrase and a 180-word cited list.
    # Temperature 0 alone is not enough — three runs still differed — but
    # temperature 0 with a fixed seed came back byte-identical.
    #
    # The seed is derived from the question, so two people asking the same
    # thing get the same answer and a different question is not pinned to the
    # same sample.
    seed = int(hashlib.sha256(" ".join(question.lower().split()).encode()).hexdigest()[:8], 16)

    async def ask():
        from app.services.ai_service import achat
        return (await achat(model=model_id, messages=messages, max_tokens=700,
                            temperature=0, seed=seed, api_key=api_key) or "").strip()

    try:
        raw = await ask()
        # One in fifteen came back empty from a healthy model in 4.9 seconds —
        # not a refusal, not an error, just nothing. Asking once more costs a
        # second and turns a dead end into an answer.
        if not raw:
            log.warning("AI Mode: empty reply for user %s, asking once more", current_user.id)
            raw = await ask()
    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.")

    # An empty answer is a failure, not an answer. It used to be stored and
    # drawn as a blank card the learner could neither read nor retry — a silent
    # failure is the worst kind, because it looks like the product working.
    if not raw.strip():
        raise HTTPException(502, "The model returned nothing. Ask again.")

    # 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)
    # Citations can be the whole of a short reply — "See [[article:7]]." with an
    # invented marker leaves an empty string once the marker is deleted.
    if not reply.strip():
        raise HTTPException(502, "The model's answer did not survive checking. Ask again.")

    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, and a rail of threads called "hi" ages worse. A greeting is not what
    # the conversation turned out to be about, so the name waits for the turn
    # that is — which is usually the very next one.
    if conversation.title == "New chat" and mode != "chat":
        conversation.title = ai_mode_service.thread_title(question)
    conversation.updated_at = datetime.utcnow()
    db.commit()
    db.refresh(answer)
    return {"message": _message_json(answer), "title": conversation.title,
            "source_count": len(sources)}