from typing import Optional

from pydantic import BaseModel

from lqs.interface.dsm.models.__common__ import (
    DataResponseModel,
    PaginationModel,
)


class AgentRun(BaseModel):
    prompt: str
    output: Optional[str]
    messages: Optional[dict | list]
    metadata: Optional[dict]
    error: Optional[str]


class AgentRunDataResponse(DataResponseModel[AgentRun]):
    pass


class AgentRunListResponse(PaginationModel[AgentRun]):
    pass


class AgentRunCreateRequest(BaseModel):
    prompt: str
    message_history: Optional[list[dict]] = None
    model_name: str = "haiku"
    thinking_enabled: bool = False
    thinking_budget_tokens: Optional[int] = 1024
    include_all_messages: bool = False


class AgentRunUpdateRequest(BaseModel):
    pass
