Livrare LOT 1 - Didi

ai_platform/modules/llm-inference/src/llm_inference/schemas.py

@@ -0,0 +1,223 @@
+"""Pydantic V2 API request/response schemas."""
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from llm_inference.types import (
+    BackendType,
+    ChatMessage,
+    Choice,
+    ModelInfo,
+    StreamChoice,
+    Usage,
+)
+
+# =============================================================================
+# Completion Schemas
+# =============================================================================
+
+
+class CompletionRequest(BaseModel):
+    """Request schema for chat completions."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    messages: list[ChatMessage] = Field(
+        min_length=1,
+        max_length=1000,
+        description="List of messages in the conversation (1-1000 messages)",
+    )
+    model: str = Field(
+        description="Model to use for completion",
+    )
+    temperature: float = Field(
+        default=0.7,
+        ge=0.0,
+        le=2.0,
+        description="Sampling temperature",
+    )
+    max_tokens: int | None = Field(
+        default=None,
+        ge=1,
+        le=1000000,
+        description="Maximum tokens to generate (1-1000000)",
+    )
+    stream: bool = Field(
+        default=False,
+        description="Enable streaming response",
+    )
+    backend: BackendType | None = Field(
+        default=None,
+        description="Backend to use (overrides default)",
+    )
+    top_p: float | None = Field(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        description="Top-p sampling parameter",
+    )
+    frequency_penalty: float | None = Field(
+        default=None,
+        ge=-2.0,
+        le=2.0,
+        description="Frequency penalty",
+    )
+    presence_penalty: float | None = Field(
+        default=None,
+        ge=-2.0,
+        le=2.0,
+        description="Presence penalty",
+    )
+    stop: list[str] | str | None = Field(
+        default=None,
+        description="Stop sequences",
+    )
+
+    @field_validator("messages")
+    @classmethod
+    def validate_messages_content(cls, v: list[ChatMessage]) -> list[ChatMessage]:
+        """Validate that messages have content (except assistant messages)."""
+        for i, msg in enumerate(v):
+            # Assistant messages can have empty content (for function calls etc)
+            if msg.role != "assistant" and msg.content is None:
+                raise ValueError(
+                    f"Message at index {i} with role '{msg.role}' cannot have empty content"
+                )
+        return v
+
+
+class CompletionResponse(BaseModel):
+    """Response schema for chat completions."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    id: str = Field(description="Unique completion ID")
+    object: str = Field(default="chat.completion")
+    created: int = Field(description="Unix timestamp of creation")
+    model: str = Field(description="Model used for completion")
+    choices: list[Choice] = Field(description="Completion choices")
+    usage: Usage | None = Field(default=None, description="Token usage")
+    backend: str = Field(description="Backend that served the request")
+
+
+class CompletionChunk(BaseModel):
+    """Streaming chunk for chat completions."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    id: str = Field(description="Unique completion ID")
+    object: str = Field(default="chat.completion.chunk")
+    created: int = Field(description="Unix timestamp of creation")
+    model: str = Field(description="Model used for completion")
+    choices: list[StreamChoice] = Field(description="Streaming choices")
+
+
+# =============================================================================
+# Text Completion Schemas (legacy OpenAI /v1/completions)
+# =============================================================================
+
+
+class TextChoice(BaseModel):
+    """One text-completion choice."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    index: int = Field(description="Choice index")
+    text: str = Field(description="Generated text")
+    finish_reason: str | None = Field(default=None, description="Finish reason")
+
+
+class TextCompletionRequest(BaseModel):
+    """Request schema for text completions (/v1/completions)."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    prompt: str | list[str] = Field(description="Prompt(s) to complete")
+    model: str | None = Field(default=None, description="Model (alias) to use")
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    max_tokens: int | None = Field(default=None, ge=1, le=1000000)
+    backend: BackendType | None = Field(
+        default=None, description="Backend to use (overrides default)"
+    )
+    top_p: float | None = Field(default=None, ge=0.0, le=1.0)
+    frequency_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
+    presence_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
+    stop: list[str] | str | None = Field(default=None, description="Stop sequences")
+
+
+class TextCompletionResponse(BaseModel):
+    """Response schema for text completions."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    id: str = Field(description="Unique completion ID")
+    object: str = Field(default="text_completion")
+    created: int = Field(description="Unix timestamp of creation")
+    model: str = Field(description="Model used for completion")
+    choices: list[TextChoice] = Field(description="Completion choices")
+    usage: Usage | None = Field(default=None, description="Token usage")
+    backend: str = Field(description="Backend that served the request")
+
+
+# =============================================================================
+# Model Management Schemas
+# =============================================================================
+
+
+class ModelListResponse(BaseModel):
+    """Response schema for listing models."""
+
+    object: str = Field(default="list")
+    data: list[ModelInfo] = Field(description="List of available models")
+
+
+class ModelLoadRequest(BaseModel):
+    """Request schema for loading a model."""
+
+    model: str = Field(description="Model identifier to load")
+    backend: BackendType = Field(description="Backend to load the model on")
+
+
+class ModelLoadResponse(BaseModel):
+    """Response schema for model load/unload operations."""
+
+    success: bool = Field(description="Whether the operation succeeded")
+    model: str = Field(description="Model identifier")
+    backend: str = Field(description="Backend the operation was performed on")
+    message: str | None = Field(default=None, description="Optional message")
+
+
+# =============================================================================
+# Health Schemas
+# =============================================================================
+
+
+class BackendHealth(BaseModel):
+    """Health status for a single backend."""
+
+    name: str = Field(description="Backend name")
+    healthy: bool = Field(description="Whether the backend is healthy")
+    message: str | None = Field(default=None, description="Optional status message")
+
+
+class HealthResponse(BaseModel):
+    """Response schema for health check."""
+
+    status: str = Field(description="Overall status: healthy, degraded, unhealthy")
+    backends: list[BackendHealth] = Field(description="Per-backend health status")
+
+
+class ReadinessResponse(BaseModel):
+    """Response schema for readiness probe."""
+
+    ready: bool = Field(default=True)
+
+
+# =============================================================================
+# Backend Schemas
+# =============================================================================
+
+
+class BackendListResponse(BaseModel):
+    """Response schema for listing backends."""
+
+    backends: list[str] = Field(description="List of available backend names")