legco_ai_assistant/backend/app/models/query.py

46 lines
839 B
Python

from typing import List, Literal, Union
from pydantic import BaseModel
from app.models.common import SourceMetadata
class QueryRequest(BaseModel):
question: str
class QueryResponse(BaseModel):
extracted_questions: List[str]
answer: str
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 CompletedEvent(BaseModel):
phase: Literal["completed"]
answer: str
sources: List[SourceMetadata]
class ErrorEvent(BaseModel):
phase: Literal["error"]
message: str
StreamingQueryEvent = Union[
DecomposedEvent, RetrievingEvent, FilteringEvent,
CompletedEvent, ErrorEvent
]