"""CoT structured output module."""

from typing import List

from pydantic import BaseModel, Field

from agential.core.llm import Response
from agential.prompting.base.output import BasePromptingOutput


class CoTStepOutput(BaseModel):
    """CoT step Pydantic output class.

    Attributes:
        thought (str): The thought of the step.
        answer (str): The answer of the step.
        thought_response (Response): The llm response of the thought.
        answer_response (Response): The llm response of the answer.
    """

    thought: str = Field(..., description="The thought of the step.")
    answer: str = Field(..., description="The answer of the step.")
    thought_response: Response = Field(
        ..., description="The llm response of the thought."
    )
    answer_response: Response = Field(
        ..., description="The llm response of the answer."
    )


class CoTOutput(BasePromptingOutput):
    """CoT Pydantic output class.

    Attributes:
        additional_info (List[List[CoTStepOutput]]): The list of list of llm responses information.
    """

    additional_info: List[List[CoTStepOutput]] = Field(
        ..., description="The list of list of llm responses information."
    )
