brief-rags-bench/app/api/v1/query.py

"""
Query API endpoints.

Отправка запросов к RAG backends в двух режимах:
1. Bench mode - batch запросы
2. Backend mode - последовательные запросы
"""

from fastapi import APIRouter, Depends, HTTPException, status
from typing import Dict, Any
from app.models.query import BenchQueryRequest, BackendQueryRequest, QueryResponse
from app.interfaces.db_api_client import DBApiClient
from app.services.rag_service import RagService
from app.dependencies import get_db_client, get_current_user
import httpx
import logging
from datetime import datetime, UTC
import uuid

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/query", tags=["query"])


def get_rag_service() -> RagService:
    """Dependency для получения RAG service."""
    return RagService()


@router.post("/bench", response_model=QueryResponse)
async def bench_query(
    request: BenchQueryRequest,
    current_user: dict = Depends(get_current_user),
    db_client: DBApiClient = Depends(get_db_client),
    rag_service: RagService = Depends(get_rag_service)
):
    """
    Отправить batch запрос к RAG backend (bench mode).

    Отправляет все вопросы одним запросом.

    Args:
        request: Запрос с окружением и списком вопросов

    Returns:
        QueryResponse: Ответ от RAG backend с metadata
    """
    user_id = current_user["user_id"]
    environment = request.environment.lower()

    
    if environment not in ['ift', 'psi', 'prod']:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid environment. Must be one of: ift, psi, prod"
        )

    try:
        
        user_settings_response = await db_client.get_user_settings(user_id)
        env_settings = user_settings_response.settings.get(environment)

        if not env_settings:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Settings not found for environment: {environment}"
            )

        
        if env_settings.apiMode != 'bench':
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Environment {environment} is not configured for bench mode"
            )

        
        request_id = str(uuid.uuid4())

        logger.info(
            f"User {user_id} sending bench query to {environment}: "
            f"{len(request.questions)} questions, request_id={request_id}"
        )

        
        env_settings_dict = env_settings.model_dump()
        response_data = await rag_service.send_bench_query(
            environment=environment,
            questions=request.questions,
            user_settings=env_settings_dict,
            request_id=request_id
        )


        return QueryResponse(
            request_id=request_id,
            timestamp=datetime.now(UTC).isoformat().replace('+00:00', 'Z'),
            environment=environment,
            response=response_data
        )

    except httpx.HTTPStatusError as e:
        logger.error(f"RAG backend error for bench query: {e}")
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail=f"RAG backend returned error: {e.response.status_code}"
        )
    except Exception as e:
        logger.error(f"Unexpected error in bench query: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Internal server error"
        )
    finally:
        await rag_service.close()


@router.post("/backend", response_model=QueryResponse)
async def backend_query(
    request: BackendQueryRequest,
    current_user: dict = Depends(get_current_user),
    db_client: DBApiClient = Depends(get_db_client),
    rag_service: RagService = Depends(get_rag_service)
):
    """
    Отправить запросы к RAG backend (backend mode).

    Отправляет вопросы по одному с возможностью сброса сессии.

    Args:
        request: Запрос с окружением, вопросами и флагом reset_session

    Returns:
        QueryResponse: Список ответов от RAG backend с metadata
    """
    user_id = current_user["user_id"]
    environment = request.environment.lower()

    
    if environment not in ['ift', 'psi', 'prod']:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid environment. Must be one of: ift, psi, prod"
        )

    try:
        
        user_settings_response = await db_client.get_user_settings(user_id)
        env_settings = user_settings_response.settings.get(environment)

        if not env_settings:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Settings not found for environment: {environment}"
            )

        
        if env_settings.apiMode != 'backend':
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Environment {environment} is not configured for backend mode"
            )

        
        request_id = str(uuid.uuid4())

        logger.info(
            f"User {user_id} sending backend query to {environment}: "
            f"{len(request.questions)} questions, request_id={request_id}, "
            f"reset_session={request.reset_session}"
        )

        
        env_settings_dict = env_settings.model_dump()
        response_data = await rag_service.send_backend_query(
            environment=environment,
            questions=request.questions,
            user_settings=env_settings_dict,
            reset_session=request.reset_session
        )


        return QueryResponse(
            request_id=request_id,
            timestamp=datetime.now(UTC).isoformat().replace('+00:00', 'Z'),
            environment=environment,
            response={"answers": response_data}
        )

    except httpx.HTTPStatusError as e:
        logger.error(f"RAG backend error for backend query: {e}")
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail=f"RAG backend returned error: {e.response.status_code}"
        )
    except Exception as e:
        logger.error(f"Unexpected error in backend query: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Internal server error"
        )
    finally:
        await rag_service.close()