| from fastapi import APIRouter, Depends, HTTPException |
| from app.core.security import get_current_user_id |
| from app.schemas.schemas import QueryRequest, QueryResponse |
| from app.services.rag_pipeline import run_rag_pipeline |
| from app.db import chat_db |
| from app.db.database import get_db |
| from app.services.ingestion_service import get_conversation_ready_document_count |
| from sqlalchemy.ext.asyncio import AsyncSession |
| from loguru import logger |
|
|
| router = APIRouter(prefix="/query", tags=["rag"]) |
|
|
|
|
| @router.post("", response_model=QueryResponse) |
| async def query( |
| request: QueryRequest, |
| user_id: str = Depends(get_current_user_id), |
| db: AsyncSession = Depends(get_db), |
| ): |
| |
| conv = await chat_db.get_conversation(user_id, request.conversation_id) |
| if not conv: |
| raise HTTPException(status_code=404, detail="Conversation not found") |
| ready_docs = await get_conversation_ready_document_count(request.conversation_id, db) |
| if ready_docs == 0: |
| await chat_db.add_message( |
| user_id=user_id, |
| conv_id=request.conversation_id, |
| role="user", |
| content=request.query, |
| ) |
| asst_msg = await chat_db.add_message( |
| user_id=user_id, |
| conv_id=request.conversation_id, |
| role="assistant", |
| content="Upload documents to this chat first", |
| ) |
| messages = await chat_db.get_messages(user_id, request.conversation_id) |
| if len(messages) <= 2: |
| await chat_db.auto_title_conversation(user_id, request.conversation_id, request.query) |
| return QueryResponse( |
| message_id=asst_msg.id, |
| conversation_id=request.conversation_id, |
| query=request.query, |
| search_query=request.query, |
| answer="Upload documents to this chat first", |
| sources=[], |
| model="none", |
| retrieval_strategy="none", |
| conversation_turns_used=0, |
| retrieval_ms=0, |
| rerank_ms=0, |
| generation_ms=0, |
| ) |
|
|
| logger.info(f"RAG query user={user_id} conv={request.conversation_id}: {request.query[:60]}...") |
| try: |
| return await run_rag_pipeline(request, user_id) |
| except Exception as exc: |
| logger.error(f"RAG pipeline error: {exc}") |
| raise HTTPException(status_code=500, detail=f"RAG pipeline failed: {str(exc)}") |
|
|
