tatlock/src/responses/schemas.py

"""
Response schemas for Responses API.

OpenAI Responses API format with support for:
- Reasoning items (thinking/reasoning summaries)
- Function call items (tool execution)
- Message items (assistant responses)
- Streaming and non-streaming modes
"""

from typing import Any, Literal

from pydantic import Field, field_validator

from src.core.models import CustomBaseModel

# ============================================================================
# Output Item Schemas (appear in response.output array)
# ============================================================================


class OutputTextContent(CustomBaseModel):
    """Text content in message output."""

    type: Literal["output_text"] = "output_text"
    text: str
    annotations: list[dict] = Field(default_factory=list)


class MessageOutputItem(CustomBaseModel):
    """
    Message item in output array.

    Represents the assistant's final response message.
    """

    type: Literal["message"] = "message"
    id: str
    role: Literal["assistant"] = "assistant"
    content: list[OutputTextContent]
    status: Literal["completed", "in_progress", "failed"] = "completed"


class ReasoningOutputItem(CustomBaseModel):
    """
    Reasoning item in output array.

    Represents the model's thinking/reasoning process.
    Displayed separately from the final answer.
    """

    type: Literal["reasoning"] = "reasoning"
    id: str
    summary: list[str]  # List of reasoning steps
    status: Literal["completed", "in_progress", "failed"] = "completed"


class FunctionCallOutputItem(CustomBaseModel):
    """
    Function call item in output array.

    Represents a tool/function that the model wants to execute.
    """

    type: Literal["function_call"] = "function_call"
    id: str
    name: str
    arguments: str  # JSON string of arguments
    status: Literal["completed", "in_progress", "failed"] = "completed"


# Union type for all output items
# Type: ignore because Pydantic handles union types specially
OutputItem = MessageOutputItem | ReasoningOutputItem | FunctionCallOutputItem


# ============================================================================
# Usage Tracking
# ============================================================================


class ResponseUsage(CustomBaseModel):
    """Token usage statistics for the response."""

    input_tokens: int
    output_tokens: int
    reasoning_tokens: int = 0
    total_tokens: int


# ============================================================================
# Request Schema
# ============================================================================


class Tool(CustomBaseModel):
    """Tool/function definition."""

    name: str
    description: str
    parameters: dict[str, Any]


class ReasoningConfig(CustomBaseModel):
    """Reasoning configuration."""

    effort: Literal["none", "minimal", "low", "medium", "high", "xhigh"] = "medium"
    summary: Literal["auto", "off"] = "auto"


class ResponseRequest(CustomBaseModel):
    """
    Request to create a response.

    OpenAI Responses API format with optional extensions.
    """

    model: str = Field(description="Model ID to use")
    input: list[dict] = Field(description="Input messages or previous responses")
    reasoning: dict[str, Any] | None = Field(
        default=None, description="Reasoning configuration: {effort: 'medium', summary: 'auto'}"
    )
    tools: list[dict] | None = Field(default=None, description="Available tools/functions")
    metadata: dict[str, Any] | None = Field(
        default=None, description="Custom metadata (e.g., conversation_id for server-side tracking)"
    )
    stream: bool = Field(default=False, description="Enable streaming mode")
    max_output_tokens: int | None = Field(default=None, description="Maximum tokens to generate")
    temperature: float = Field(default=1.0, ge=0.0, le=2.0, description="Sampling temperature")
    stop: list[str] | None = Field(default=None, description="Stop sequences")
    user: str | None = Field(
        default=None, description="Unique identifier for end-user (OpenAI standard)"
    )

    @field_validator("reasoning")
    @classmethod
    def validate_reasoning(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
        """
        Validate reasoning configuration.

        Checks:
        - effort must be valid level (none, minimal, low, medium, high, xhigh)
        - summary must be 'auto' or 'off'
        """
        if v is not None:
            if "effort" in v:
                allowed_efforts = ["none", "minimal", "low", "medium", "high", "xhigh"]
                if v["effort"] not in allowed_efforts:
                    raise ValueError(
                        f"reasoning.effort must be one of {allowed_efforts}, got '{v['effort']}'"
                    )
            if "summary" in v:
                allowed_summaries = ["auto", "off"]
                if v["summary"] not in allowed_summaries:
                    raise ValueError(
                        f"reasoning.summary must be one of {allowed_summaries}, got '{v['summary']}'"
                    )
        return v

    @field_validator("max_output_tokens")
    @classmethod
    def validate_max_output_tokens(cls, v: int | None) -> int | None:
        """
        Validate max_output_tokens.

        Must be positive if provided.
        """
        if v is not None and v <= 0:
            raise ValueError(f"max_output_tokens must be positive, got {v}")
        return v

    @field_validator("stop")
    @classmethod
    def validate_stop_sequences(cls, v: list[str] | None) -> list[str] | None:
        """
        Validate stop sequences.

        Checks:
        - Maximum 4 stop sequences
        - Each must be non-empty string
        """
        if v is not None:
            if len(v) > 4:
                raise ValueError(f"Maximum 4 stop sequences allowed, got {len(v)}")
            for seq in v:
                if not seq or not isinstance(seq, str):
                    raise ValueError("Stop sequences must be non-empty strings")
        return v


# ============================================================================
# Response Schema
# ============================================================================


class Response(CustomBaseModel):
    """
    Complete response object.

    Contains output array with reasoning, function calls, and messages.
    """

    id: str = Field(description="Unique response ID")
    object: Literal["response"] = "response"
    created_at: int = Field(description="Unix timestamp")
    model: str = Field(description="Model used")
    status: Literal["completed", "in_progress", "failed", "cancelled"]
    output: list[OutputItem] = Field(description="Output items (reasoning, function_call, message)")
    usage: ResponseUsage = Field(description="Token usage statistics")


# ============================================================================
# Error Schema
# ============================================================================


class ErrorDetail(CustomBaseModel):
    """Error detail object."""

    type: str
    message: str
    code: int | None = None


class ErrorResponse(CustomBaseModel):
    """Error response format."""

    error: ErrorDetail