SupplyAgent/supply_agent/types.py

from __future__ import annotations

from collections.abc import Callable, Sequence
from enum import StrEnum
from typing import Any, TypeAlias

from pydantic import BaseModel, Field


class Role(StrEnum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


class Message(BaseModel):
    """A single message in the conversation."""

    role: Role
    content: str | None = None
    tool_calls: list[ToolCall] | None = None
    tool_call_id: str | None = None
    name: str | None = None
    # Model reasoning / chain-of-thought text (OpenRouter); not sent back in API history.
    reasoning: str | None = None

    def to_api_dict(self) -> dict[str, Any]:
        """Convert to OpenAI-compatible API format."""
        data: dict[str, Any] = {"role": self.role.value}
        if self.content is not None:
            data["content"] = self.content
        if self.tool_calls:
            data["tool_calls"] = [tc.to_api_dict() for tc in self.tool_calls]
        if self.tool_call_id:
            data["tool_call_id"] = self.tool_call_id
        if self.name:
            data["name"] = self.name
        return data


CompletionGuard: TypeAlias = Callable[
    [Message, Sequence[Message]],
    str | None,
]
"""Completion callback.

Return ``None`` to accept an assistant response without tool calls as the final
answer. Return feedback text to reject completion and continue the agent loop;
the feedback is appended as a user message before the next iteration. The guard
does not run for the forced final answer after ``max_iterations`` is reached.
"""


class ToolCall(BaseModel):
    """A tool invocation requested by the model."""

    id: str
    name: str
    arguments: str  # JSON string

    def to_api_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "type": "function",
            "function": {
                "name": self.name,
                "arguments": self.arguments,
            },
        }


class ToolDefinition(BaseModel):
    """OpenAI-compatible tool schema."""

    name: str
    description: str
    parameters: dict[str, Any] = Field(default_factory=lambda: {
        "type": "object",
        "properties": {},
        "required": [],
    })

    def to_api_dict(self) -> dict[str, Any]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


class ToolResult(BaseModel):
    """Result of executing a tool."""

    tool_call_id: str
    name: str
    content: str
    is_error: bool = False


class SkillInfo(BaseModel):
    """Metadata and content for a loaded skill."""

    name: str
    description: str
    content: str
    path: str


class AgentEventType(StrEnum):
    """Events emitted during agent execution."""

    THINKING = "thinking"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"
    MESSAGE = "message"
    SKILL_LOADED = "skill_loaded"
    DONE = "done"
    ERROR = "error"


class AgentEvent(BaseModel):
    """A streaming event from the agent loop."""

    type: AgentEventType
    data: dict[str, Any] = Field(default_factory=dict)


class AgentResult(BaseModel):
    """Final result of an agent run."""

    content: str
    messages: list[Message]
    iterations: int
    tool_calls_made: int
    skills_used: list[str] = Field(default_factory=list)