"""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)}