223 lines
7.4 KiB
Python
223 lines
7.4 KiB
Python
"""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")
|