from __future__ import annotations

from typing import ClassVar

from pydantic import BaseModel, ConfigDict

from ..agent import AgentType, ToolAgentOutput, WorkerAgentOutput


class UserMessageAttachments(BaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

    bucket: str
    path: str
    mime_type: str


class AgentChatMessageMetadata(BaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

    agent_type: AgentType | None = None
    user_message_attachments: UserMessageAttachments | None = None
    tool_agent_output: ToolAgentOutput | None = None
    worker_agent_output: WorkerAgentOutput | None = None