src/api/schemas.py

"""
Frox AI — API Schemas

OpenAI-compatible request/response shapes. Matching this contract
precisely (not "close enough") is what lets Open WebUI, most desktop
AI clients, and any future CLI/Android client point at this server
via a base-URL setting with zero code changes on their end.
"""
from __future__ import annotations

import time
import uuid
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field


class ToolCallFunction(BaseModel):
    name: str
    arguments: str    # JSON-encoded string, per OpenAI convention — not a nested object


class ToolCall(BaseModel):
    id: str
    type: Literal["function"] = "function"
    function: ToolCallFunction


class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant", "tool"]
    content: Optional[str] = None
    name: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None
    tool_call_id: Optional[str] = None


class ChatCompletionRequest(BaseModel):
    model: str = "classic"                       # maps to config/family/<name>.py
    messages: List[ChatMessage]
    stream: bool = False
    max_tokens: int = Field(default=1024, ge=1, le=32768)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
    top_k: int = Field(default=50, ge=0)
    repetition_penalty: float = Field(default=1.1, ge=0.1, le=2.0)
    user: Optional[str] = None                    # OpenAI convention: caller-supplied user id
    session_id: Optional[str] = None               # Frox extension: enables session KV-cache reuse
    tools: Optional[List[Dict[str, Any]]] = None    # OpenAI-style function schemas


class ChatCompletionChoice(BaseModel):
    index: int = 0
    message: ChatMessage
    finish_reason: Literal["stop", "length", "tool_calls"] = "stop"


class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class ChatCompletionResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:24]}")
    object: Literal["chat.completion"] = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[ChatCompletionChoice]
    usage: Usage
    frox_trace: Optional[List[Dict[str, Any]]] = None   # populated for model="conductor" responses


class ChatCompletionChunkDelta(BaseModel):
    role: Optional[Literal["assistant"]] = None
    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None


class ChatCompletionChunkChoice(BaseModel):
    index: int = 0
    delta: ChatCompletionChunkDelta
    finish_reason: Optional[Literal["stop", "length", "tool_calls"]] = None


class ChatCompletionChunk(BaseModel):
    id: str
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[ChatCompletionChunkChoice]


class ModelInfo(BaseModel):
    id: str
    object: Literal["model"] = "model"
    created: int = Field(default_factory=lambda: int(time.time()))
    owned_by: str = "frox-ai"


class ModelListResponse(BaseModel):
    object: Literal["list"] = "list"
    data: List[ModelInfo]


class ErrorDetail(BaseModel):
    message: str
    type: str = "invalid_request_error"
    code: Optional[str] = None
    param: Optional[str] = None


class ErrorResponse(BaseModel):
    error: ErrorDetail


class ToolExecuteRequest(BaseModel):
    tool: str
    args: Dict[str, Any] = Field(default_factory=dict)
    user_id: str = "default"
    session_id: Optional[str] = None
    plan: str = "free"


class HealthResponse(BaseModel):
    status: Literal["ok", "loading", "error"]
    model: Optional[str] = None
    device: Optional[str] = None
    version: str = "1.1.0"