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
]