from typing import Any

from pydantic import BaseModel, Field

from aiteamwork.llm_context import LLMContext


class LLMHumanInTheLoopContext[ToolArgsType: BaseModel, RuntimeContextType: BaseModel](LLMContext[RuntimeContextType]):
    """Context for the instructions callback in LLMAgent."""

    tool_args: ToolArgsType = Field()
    """Arguments for the tool call."""


class LLMHumanInTheLoopResult[ArtifactType: BaseModel | None](BaseModel):
    """Result of the instructions callback in LLMAgent."""

    artifact: ArtifactType = Field()
    """Artifact generated by the confirmation."""

    class Config:
        arbitrary_types_allowed = True


class LLMHumanInTheLoopArtifacts(BaseModel):
    """Artifacts produced for the human-in-the-loop, derived from the tool arguments"""

    artifacts: list[BaseModel] = Field()
    """Artifacts produced for the human-in-the-loop, derived from the tool arguments"""

    class Config:
        arbitrary_types_allowed = True


class LLMHumanInTheLoopHumanInputs(BaseModel):
    """Resulting human inputs from the human-in-the-loop, passed back to the tool in the context.human_input field"""

    inputs: list[Any] = Field()
    """Resulting human inputs from the human-in-the-loop, passed back to the tool in the context.human_input field"""

    class Config:
        arbitrary_types_allowed = True
