Types

class GenerateResponseWrapper(GenerateResponse):
    """A wrapper around GenerateResponse providing utility methods.

    Extends the base `GenerateResponse` with cached properties (`text`, `output`,
    `messages`, `tool_requests`) and methods for validation (`assert_valid`,
    `assert_valid_schema`). It also handles optional message/chunk parsing.
    """

    message_parser: MessageParser | None = Field(exclude=True)
    message: MessageWrapper = None

    def __init__(
        self,
        response: GenerateResponse,
        request: GenerateRequest,
        message_parser: MessageParser | None = None,
    ):
        """Initializes a GenerateResponseWrapper instance.

        Args:
            response: The original GenerateResponse object.
            request: The GenerateRequest object associated with the response.
            message_parser: An optional function to parse the output from the message.
        """
        super().__init__(
            message=MessageWrapper(response.message)
            if not isinstance(response.message, MessageWrapper)
            else response.message,
            finish_reason=response.finish_reason,
            finish_message=response.finish_message,
            latency_ms=response.latency_ms,
            usage=response.usage if response.usage is not None else GenerationUsage(),
            custom=response.custom if response.custom is not None else {},
            request=request,
            candidates=response.candidates,
            message_parser=message_parser,
        )

    def assert_valid(self):
        """Validates the basic structure of the response.

        Note: This method is currently a placeholder (TODO).

        Raises:
            AssertionError: If the response structure is considered invalid.
        """
        # TODO: implement
        pass

    def assert_valid_schema(self):
        """Validates that the response message conforms to any specified output schema.

        Note: This method is currently a placeholder (TODO).

        Raises:
            AssertionError: If the response message does not conform to the schema.
        """
        # TODO: implement
        pass

    @cached_property
    def text(self) -> str:
        """Returns all text parts of the response joined into a single string.

        Returns:
            str: The combined text content from the response.
        """
        return self.message.text

    @cached_property
    def output(self) -> Any:
        """Parses out JSON data from the text parts of the response.

        Returns:
            Any: The parsed JSON data from the response.
        """
        if self.message_parser:
            return self.message_parser(self.message)
        return extract_json(self.text)

    @cached_property
    def messages(self) -> list[Message]:
        """Returns all messages of the response, including request messages as a list.

        Returns:
            list[Message]: list of messages.
        """
        return [
            *(self.request.messages if self.request else []),
            self.message._original_message,
        ]

    @cached_property
    def tool_requests(self) -> list[ToolRequestPart]:
        """Returns all tool request parts of the response as a list.

        Returns:
            list[ToolRequestPart]: list of tool requests present in this response.
        """
        return self.message.tool_requests

    @cached_property
    def interrupts(self) -> list[ToolRequestPart]:
        """Returns all interrupted tool request parts of the response as a list.

        Returns:
            list[ToolRequestPart]: list of interrupted tool requests.
        """
        return self.message.interrupts

class GenerateRequest(BaseModel):
    """Model for generaterequest data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    messages: list[Message]
    config: Any | None = None
    tools: list[ToolDefinition] | None = None
    tool_choice: ToolChoice | None = Field(None, alias='toolChoice')
    output: OutputConfig | None = None
    docs: list[DocumentData] | None = None
    candidates: float | None = None

class GenerateResponse(BaseModel):
    """Model for generateresponse data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    message: Message | None = None
    finish_reason: FinishReason | None = Field(None, alias='finishReason')
    finish_message: str | None = Field(None, alias='finishMessage')
    latency_ms: float | None = Field(None, alias='latencyMs')
    usage: GenerationUsage | None = None
    custom: Any | None = None
    raw: Any | None = None
    request: GenerateRequest | None = None
    candidates: list[Candidate] | None = None

class MessageWrapper(Message):
    """A wrapper around the base Message type providing utility methods.

    This class extends the standard `Message` by adding convenient cached properties
    like `text` (for concatenated text content) and `tool_requests`.
    It stores the original message in the `_original_message` attribute.
    """

    def __init__(
        self,
        message: Message,
    ):
        """Initializes the MessageWrapper.

        Args:
            message: The original Message object to wrap.
        """
        super().__init__(
            role=message.role,
            content=message.content,
            metadata=message.metadata,
        )
        self._original_message = message

    @cached_property
    def text(self) -> str:
        """Returns all text parts of the current chunk joined into a single string.

        Returns:
            str: The combined text content from the current chunk.
        """
        return text_from_message(self)

    @cached_property
    def tool_requests(self) -> list[ToolRequestPart]:
        """Returns all tool request parts of the response as a list.

        Returns:
            list[ToolRequestPart]: list of tool requests present in this response.
        """
        return [p.root for p in self.content if isinstance(p.root, ToolRequestPart)]

    @cached_property
    def interrupts(self) -> list[ToolRequestPart]:
        """Returns all interrupted tool request parts of the message as a list.

        Returns:
            list[ToolRequestPart]: list of interrupted tool requests.
        """
        return [p for p in self.tool_requests if p.metadata and p.metadata.root.get('interrupt')]

class MessageWrapper(Message):
    """A wrapper around the base Message type providing utility methods.

    This class extends the standard `Message` by adding convenient cached properties
    like `text` (for concatenated text content) and `tool_requests`.
    It stores the original message in the `_original_message` attribute.
    """

    def __init__(
        self,
        message: Message,
    ):
        """Initializes the MessageWrapper.

        Args:
            message: The original Message object to wrap.
        """
        super().__init__(
            role=message.role,
            content=message.content,
            metadata=message.metadata,
        )
        self._original_message = message

    @cached_property
    def text(self) -> str:
        """Returns all text parts of the current chunk joined into a single string.

        Returns:
            str: The combined text content from the current chunk.
        """
        return text_from_message(self)

    @cached_property
    def tool_requests(self) -> list[ToolRequestPart]:
        """Returns all tool request parts of the response as a list.

        Returns:
            list[ToolRequestPart]: list of tool requests present in this response.
        """
        return [p.root for p in self.content if isinstance(p.root, ToolRequestPart)]

    @cached_property
    def interrupts(self) -> list[ToolRequestPart]:
        """Returns all interrupted tool request parts of the message as a list.

        Returns:
            list[ToolRequestPart]: list of interrupted tool requests.
        """
        return [p for p in self.tool_requests if p.metadata and p.metadata.root.get('interrupt')]

class Message(BaseModel):
    """Model for message data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    role: Role
    content: list[Part]
    metadata: dict[str, Any] | None = None

class Part(RootModel[TextPart | MediaPart | ToolRequestPart | ToolResponsePart | DataPart | CustomPart]):
    """Root model for part."""

    root: TextPart | MediaPart | ToolRequestPart | ToolResponsePart | DataPart | CustomPart

class TextPart(BaseModel):
    """Model for textpart data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    text: str
    media: MediaModel | None = None
    tool_request: ToolRequestModel | None = Field(None, alias='toolRequest')
    tool_response: ToolResponseModel | None = Field(None, alias='toolResponse')
    data: Data | None = None
    metadata: Metadata | None = None
    custom: Custom | None = None

class MediaPart(BaseModel):
    """Model for mediapart data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    text: Text | None = None
    media: Media
    tool_request: ToolRequestModel | None = Field(None, alias='toolRequest')
    tool_response: ToolResponseModel | None = Field(None, alias='toolResponse')
    data: Data | None = None
    metadata: Metadata | None = None
    custom: Custom | None = None

class Media(BaseModel):
    """Model for media data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    content_type: str | None = Field(None, alias='contentType')
    url: str

class ToolRequestPart(BaseModel):
    """Model for toolrequestpart data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    text: Text | None = None
    media: MediaModel | None = None
    tool_request: ToolRequest = Field(..., alias='toolRequest')
    tool_response: ToolResponseModel | None = Field(None, alias='toolResponse')
    data: Data | None = None
    metadata: Metadata | None = None
    custom: Custom | None = None

class ToolRequest(BaseModel):
    """Model for toolrequest data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    ref: str | None = None
    name: str
    input: Any | None = None

class ToolResponsePart(BaseModel):
    """Model for toolresponsepart data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    text: Text | None = None
    media: MediaModel | None = None
    tool_request: ToolRequestModel | None = Field(None, alias='toolRequest')
    tool_response: ToolResponse = Field(..., alias='toolResponse')
    data: Data | None = None
    metadata: Metadata | None = None
    custom: Custom | None = None

class ToolResponse(BaseModel):
    """Model for toolresponse data."""

    model_config = ConfigDict(extra='forbid', populate_by_name=True)
    ref: str | None = None
    name: str
    output: Any | None = None