brief-rags-bench/app/models/query.py

"""Query request/response Pydantic models."""

from typing import Any
from pydantic import BaseModel


class QuestionRequest(BaseModel):
    """Single question for batch query."""

    body: str
    with_docs: bool = True


class BenchQueryRequest(BaseModel):
    """Bench mode query request."""

    environment: str  # ift, psi, prod
    questions: list[QuestionRequest]


class BackendQueryRequest(BaseModel):
    """Backend mode query request (one-by-one)."""

    environment: str  # ift, psi, prod
    questions: list[QuestionRequest]
    reset_session: bool = True


class QueryResponse(BaseModel):
    """Query response with metadata."""

    request_id: str
    timestamp: str
    environment: str
    response: dict