DocPilot/backend/app/api/chat.py

from fastapi import (
    APIRouter,
    Depends,
    HTTPException,
)

from groq import AuthenticationError

from sqlalchemy.orm import Session

from DocPilot.backend.app.services.rag import (
    ask_question,
)

from DocPilot.backend.app.core.dependencies import (
    get_current_user,
)

from DocPilot.backend.app.db.session import get_db

from DocPilot.backend.app.schemas.chat import ChatRequest

from DocPilot.backend.app.models.chat import (
    ChatSession,
    ChatMessage,
)

router = APIRouter()


@router.post("/ask")
def ask(
    query: ChatRequest,
    current_user=Depends(get_current_user),
    db: Session = Depends(get_db),
):

    session_id = query.session_id

    if not session_id:

        session = ChatSession(
            owner_id=current_user.id,
            title=query.question.strip()[:40],
            mode=query.mode,
        )

        db.add(session)

        db.commit()

        db.refresh(session)

        session_id = session.id

    user_message = ChatMessage(
        session_id=session_id,
        role="user",
        content=query.question,
    )

    db.add(user_message)

    db.commit()

    try:
        rag_response = ask_question(
            question=query.question,
            user_id=current_user.id,
            source=query.source,
            model_name=query.model_name,
            retrieval_strategy=query.retrieval_strategy,
            reranker=query.reranker,
            enhancements=query.enhancements,
            mode=query.mode,
        )
    except AuthenticationError as exc:
        raise HTTPException(
            status_code=502,
            detail="Groq authentication failed. Check GROQ_API_KEY in .env and restart the server.",
        ) from exc

    assistant_message = ChatMessage(
        session_id=session_id,
        role="assistant",
        content=rag_response["answer"],
    )

    db.add(assistant_message)

    db.commit()

    return {
        "session_id": session_id,
        "answer": rag_response["answer"],
        "sources": rag_response["sources"],
        "trace_id": rag_response.get("trace_id"),
    }