from collections.abc import Iterator

from pydantic import BaseModel


class LLMOutput(BaseModel):
    responses: list[str]
    prompt_tokens: list[int]
    completion_tokens: list[int]

    def __getitem__(self, item) -> tuple[str, int, int]:
        return self.responses[item], self.prompt_tokens[item], self.completion_tokens[item]

    def iterate_elements(self) -> Iterator[tuple[str, int, int]]:
        for i in range(len(self.responses)):
            yield self[i]

    @classmethod
    def from_list(cls, inputs: list[tuple[str, int, int]]) -> "LLMOutput":
        return cls(
            responses=[i[0] for i in inputs],
            prompt_tokens=[i[1] for i in inputs],
            completion_tokens=[i[2] for i in inputs],
        )
