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 Literal, Any
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  # type: ignore


# ============================================================================
# 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"
    )

    @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