models/openai.py

"""
Nancy — OpenAI-compatible Pydantic v2 schemas.

These models exactly match the OpenAI Chat Completions API response format
so that the official ``openai`` Python SDK works seamlessly.
"""

from __future__ import annotations

import time
import uuid
from typing import Literal

from pydantic import BaseModel, Field


# ── Helpers ───────────────────────────────────────────────────────────────────


def _chatcmpl_id() -> str:
    """Generate an OpenAI-style completion ID."""
    return f"chatcmpl-{uuid.uuid4().hex[:29]}"


def _unix_ts() -> int:
    """Current UTC Unix timestamp."""
    return int(time.time())


# ── Request Models ────────────────────────────────────────────────────────────


class ChatMessage(BaseModel):
    """A single chat message in the OpenAI format."""

    role: Literal["system", "user", "assistant", "function", "tool"] = Field(
        ..., description="The role of the message author."
    )
    content: str | None = Field(
        default=None, description="The text content of the message."
    )
    name: str | None = Field(
        default=None, description="Optional name for the message author."
    )
    tool_call_id: str | None = Field(
        default=None, description="Tool call that this message is responding to."
    )


class ChatCompletionRequest(BaseModel):
    """
    Incoming request body for ``POST /v1/chat/completions``.

    Mirrors the subset of OpenAI fields that Nancy supports.
    """

    model: str = Field(
        ..., description="Model/provider name, e.g. 'chatgpt', 'gemini'."
    )
    messages: list[ChatMessage] = Field(
        ..., min_length=1, description="Conversation messages."
    )
    stream: bool = Field(
        default=False, description="Whether to stream the response via SSE."
    )
    temperature: float | None = Field(
        default=None, ge=0.0, le=2.0, description="Sampling temperature."
    )
    max_tokens: int | None = Field(
        default=None, ge=1, description="Maximum tokens to generate."
    )
    top_p: float | None = Field(
        default=None, ge=0.0, le=1.0, description="Nucleus sampling parameter."
    )
    stop: str | list[str] | None = Field(
        default=None, description="Stop sequences."
    )
    user: str | None = Field(
        default=None, description="End-user identifier for abuse tracking."
    )
    tools: list[dict] | None = Field(
        default=None, description="A list of tools the model may call."
    )
    tool_choice: str | dict | None = Field(
        default=None, description="Controls which (if any) tool is called."
    )


# ── Response Models — Non-streaming ──────────────────────────────────────────


class UsageInfo(BaseModel):
    """Token usage statistics."""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChoiceMessage(BaseModel):
    """The assistant's response message in a non-streaming completion."""

    role: Literal["assistant"] = "assistant"
    content: str | None = ""
    tool_calls: list[dict] | None = Field(
        default=None, description="The tool calls generated by the model."
    )


class Choice(BaseModel):
    """A single choice in a non-streaming completion response."""

    index: int = 0
    message: ChoiceMessage = Field(default_factory=ChoiceMessage)
    finish_reason: Literal["stop", "length", "content_filter", "tool_calls"] | None = None


class ChatCompletionResponse(BaseModel):
    """
    Non-streaming response for ``POST /v1/chat/completions``.

    Matches ``openai.types.chat.ChatCompletion``.
    """

    id: str = Field(default_factory=_chatcmpl_id)
    object: Literal["chat.completion"] = "chat.completion"
    created: int = Field(default_factory=_unix_ts)
    model: str = ""
    choices: list[Choice] = Field(default_factory=lambda: [Choice()])
    usage: UsageInfo = Field(default_factory=UsageInfo)
    system_fingerprint: str | None = None

    @classmethod
    def from_content(
        cls,
        content: str,
        model: str,
        finish_reason: str = "stop",
    ) -> ChatCompletionResponse:
        """Build a complete response from a single content string."""
        return cls(
            model=model,
            choices=[
                Choice(
                    index=0,
                    message=ChoiceMessage(content=content),
                    finish_reason=finish_reason,  # type: ignore[arg-type]
                )
            ],
            usage=UsageInfo(
                prompt_tokens=0,
                completion_tokens=len(content.split()),
                total_tokens=len(content.split()),
            ),
        )


# ── Response Models — Streaming (SSE chunks) ─────────────────────────────────


class DeltaContent(BaseModel):
    """
    Delta object inside a streaming chunk.

    On the first chunk, ``role`` is set to ``"assistant"`` with no content.
    On subsequent chunks, ``content`` carries the text fragment.
    On the final chunk, both may be absent (empty delta).
    """

    role: Literal["assistant"] | None = None
    content: str | None = None
    tool_calls: list[dict] | None = Field(
        default=None, description="The tool calls generated by the model."
    )


class StreamChoice(BaseModel):
    """A single choice in a streaming chunk."""

    index: int = 0
    delta: DeltaContent = Field(default_factory=DeltaContent)
    finish_reason: Literal["stop", "length", "content_filter", "tool_calls"] | None = None


class ChatCompletionChunk(BaseModel):
    """
    A single SSE chunk for streaming ``POST /v1/chat/completions``.

    Matches ``openai.types.chat.ChatCompletionChunk``.
    """

    id: str = Field(default_factory=_chatcmpl_id)
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int = Field(default_factory=_unix_ts)
    model: str = ""
    choices: list[StreamChoice] = Field(default_factory=lambda: [StreamChoice()])
    system_fingerprint: str | None = None

    def to_sse_data(self) -> str:
        """Serialize to the JSON string used in ``data: ...`` SSE frames."""
        return self.model_dump_json(exclude_none=False)

    # ── Convenience factories ─────────────────────────────────────────

    @classmethod
    def first_chunk(cls, completion_id: str, model: str) -> ChatCompletionChunk:
        """Role-only opening chunk (no content)."""
        return cls(
            id=completion_id,
            model=model,
            choices=[
                StreamChoice(
                    delta=DeltaContent(role="assistant"),
                    finish_reason=None,
                )
            ],
        )

    @classmethod
    def content_chunk(
        cls, completion_id: str, model: str, content: str
    ) -> ChatCompletionChunk:
        """A chunk carrying a text fragment."""
        return cls(
            id=completion_id,
            model=model,
            choices=[
                StreamChoice(
                    delta=DeltaContent(content=content),
                    finish_reason=None,
                )
            ],
        )

    @classmethod
    def final_chunk(
        cls,
        completion_id: str,
        model: str,
        finish_reason: str = "stop",
    ) -> ChatCompletionChunk:
        """Terminal chunk with ``finish_reason`` and empty delta."""
        return cls(
            id=completion_id,
            model=model,
            choices=[
                StreamChoice(
                    delta=DeltaContent(),
                    finish_reason=finish_reason,  # type: ignore[arg-type]
                )
            ],
        )


# ── /v1/models response ──────────────────────────────────────────────────────


class ModelInfo(BaseModel):
    """A single model entry returned by ``GET /v1/models``."""

    id: str
    object: Literal["model"] = "model"
    created: int = Field(default_factory=_unix_ts)
    owned_by: str = "nancy"


class ModelListResponse(BaseModel):
    """Response body for ``GET /v1/models``."""

    object: Literal["list"] = "list"
    data: list[ModelInfo] = Field(default_factory=list)


# ── Error response ────────────────────────────────────────────────────────────


class ErrorDetail(BaseModel):
    """OpenAI-style error detail."""

    message: str
    type: str = "invalid_request_error"
    param: str | None = None
    code: str | None = None


class ErrorResponse(BaseModel):
    """OpenAI-style error envelope."""

    error: ErrorDetail