from typing import Dict, List, Optional

from pydantic import Field

from mirix.schemas.message import Message
from mirix.schemas.mirix_base import MirixBase


class StepBase(MirixBase):
    __id_prefix__ = "step"


class Step(StepBase):
    id: str = Field(..., description="The id of the step. Assigned by the database.")
    origin: Optional[str] = Field(
        None, description="The surface that this agent step was initiated from."
    )
    organization_id: Optional[str] = Field(
        None,
        description="The unique identifier of the organization associated with the step.",
    )
    provider_name: Optional[str] = Field(
        None, description="The name of the provider used for this step."
    )
    model: Optional[str] = Field(
        None, description="The name of the model used for this step."
    )
    context_window_limit: Optional[int] = Field(
        None, description="The context window limit configured for this step."
    )
    completion_tokens: Optional[int] = Field(
        None,
        description="The number of tokens generated by the agent during this step.",
    )
    prompt_tokens: Optional[int] = Field(
        None, description="The number of tokens in the prompt during this step."
    )
    total_tokens: Optional[int] = Field(
        None,
        description="The total number of tokens processed by the agent during this step.",
    )
    completion_tokens_details: Optional[Dict] = Field(
        None, description="Metadata for the agent."
    )

    tags: List[str] = Field([], description="Metadata tags.")
    tid: Optional[str] = Field(
        None,
        description="The unique identifier of the transaction that processed this step.",
    )
    messages: List[Message] = Field(
        [], description="The messages generated during this step."
    )
