72 lines
1.5 KiB
Python
72 lines
1.5 KiB
Python
from typing import List, Literal, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from app.models.common import SourceMetadata
|
|
|
|
|
|
class QueryRequest(BaseModel):
|
|
question: str
|
|
|
|
|
|
class SubQuestionSources(BaseModel):
|
|
sub_question_index: int
|
|
sub_question_text: str
|
|
sources: List[SourceMetadata]
|
|
|
|
|
|
class SubQuestionResult(BaseModel):
|
|
sub_question_index: int
|
|
sub_question_text: str
|
|
answer_section: str
|
|
|
|
|
|
class QueryResponse(BaseModel):
|
|
extracted_questions: List[str]
|
|
answer: str
|
|
sub_question_sources: List[SubQuestionSources] = []
|
|
sources: List[SourceMetadata] = []
|
|
|
|
|
|
class DecomposedEvent(BaseModel):
|
|
phase: Literal["decomposed"]
|
|
extracted_questions: List[str]
|
|
|
|
|
|
class RetrievingEvent(BaseModel):
|
|
phase: Literal["retrieving"]
|
|
|
|
|
|
class FilteringEvent(BaseModel):
|
|
phase: Literal["filtering"]
|
|
|
|
|
|
class GeneratingEvent(BaseModel):
|
|
phase: Literal["generating"]
|
|
|
|
|
|
class GeneratingSubquestionEvent(BaseModel):
|
|
phase: Literal["generating_subquestion"]
|
|
sub_question_index: int
|
|
sub_question_text: str
|
|
answer_section: str = ""
|
|
|
|
|
|
class CompletedEvent(BaseModel):
|
|
phase: Literal["completed"]
|
|
answer: str
|
|
sub_question_sources: List[SubQuestionSources] = []
|
|
sources: List[SourceMetadata] = []
|
|
|
|
|
|
class ErrorEvent(BaseModel):
|
|
phase: Literal["error"]
|
|
message: str
|
|
|
|
|
|
StreamingQueryEvent = Union[
|
|
DecomposedEvent, RetrievingEvent, FilteringEvent,
|
|
GeneratingEvent, GeneratingSubquestionEvent,
|
|
CompletedEvent, ErrorEvent
|
|
]
|