68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
"""Query request/response Pydantic models."""
|
|
|
|
from typing import Any
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
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 Docs(BaseModel):
|
|
"""Documents from RAG."""
|
|
|
|
research: list
|
|
analytical_hub: list
|
|
|
|
|
|
class RagResponse(BaseModel):
|
|
"""Ответ от RAG на вопрос пользователя."""
|
|
|
|
body_research: str = Field(description="Текст ответа от Research на вопрос")
|
|
body_analytical_hub: str = Field(description="Текст ответа от Analytical Hub на вопрос")
|
|
docs_from_vectorstore: Docs | None = None
|
|
docs_to_llm: Docs | None = None
|
|
|
|
|
|
class RagResponseBench(RagResponse):
|
|
"""Ответ на вопрос + время обработки именно этого вопроса."""
|
|
|
|
processing_time_sec: float = Field(
|
|
description="Время обработки запроса в секундах",
|
|
ge=0,
|
|
)
|
|
question: str = Field(description="Исходный вопрос")
|
|
|
|
|
|
class RagResponseBenchList(BaseModel):
|
|
"""Список ответов RAG в bench режиме."""
|
|
|
|
answers: list[RagResponseBench]
|
|
|
|
|
|
class QueryResponse(BaseModel):
|
|
"""Query response with metadata."""
|
|
|
|
request_id: str
|
|
timestamp: str
|
|
environment: str
|
|
response: RagResponseBenchList | dict | list # RagResponseBenchList для bench, dict/list для backend
|