# This file was auto-generated by Fern from our API Definition.

from __future__ import annotations
from ..core.unchecked_base_model import UncheckedBaseModel
import typing
from ..core.pydantic_utilities import IS_PYDANTIC_V2
import pydantic
from .chat_search_query import ChatSearchQuery
from .chat_search_result import ChatSearchResult
from .chat_document import ChatDocument
from .chat_citation import ChatCitation
from .tool_call import ToolCall
from .chat_stream_end_event_finish_reason import ChatStreamEndEventFinishReason
from .non_streamed_chat_response import NonStreamedChatResponse
from .tool_call_delta import ToolCallDelta
import typing_extensions
from ..core.unchecked_base_model import UnionMetadata


class StreamStartStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["stream-start"] = "stream-start"
    generation_id: str

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class SearchQueriesGenerationStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["search-queries-generation"] = "search-queries-generation"
    search_queries: typing.List[ChatSearchQuery]

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class SearchResultsStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["search-results"] = "search-results"
    search_results: typing.Optional[typing.List[ChatSearchResult]] = None
    documents: typing.Optional[typing.List[ChatDocument]] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class TextGenerationStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["text-generation"] = "text-generation"
    text: str

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class CitationGenerationStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["citation-generation"] = "citation-generation"
    citations: typing.List[ChatCitation]

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class ToolCallsGenerationStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["tool-calls-generation"] = "tool-calls-generation"
    text: typing.Optional[str] = None
    tool_calls: typing.List[ToolCall]

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class StreamEndStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["stream-end"] = "stream-end"
    finish_reason: ChatStreamEndEventFinishReason
    response: NonStreamedChatResponse

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class ToolCallsChunkStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["tool-calls-chunk"] = "tool-calls-chunk"
    tool_call_delta: ToolCallDelta
    text: typing.Optional[str] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class DebugStreamedChatResponse(UncheckedBaseModel):
    """
    StreamedChatResponse is returned in streaming mode (specified with `stream=True` in the request).
    """

    event_type: typing.Literal["debug"] = "debug"
    prompt: typing.Optional[str] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


StreamedChatResponse = typing_extensions.Annotated[
    typing.Union[
        StreamStartStreamedChatResponse,
        SearchQueriesGenerationStreamedChatResponse,
        SearchResultsStreamedChatResponse,
        TextGenerationStreamedChatResponse,
        CitationGenerationStreamedChatResponse,
        ToolCallsGenerationStreamedChatResponse,
        StreamEndStreamedChatResponse,
        ToolCallsChunkStreamedChatResponse,
        DebugStreamedChatResponse,
    ],
    UnionMetadata(discriminant="event_type"),
]
