# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/llm_schemas.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
LLM Settings Pydantic Schemas.
This module provides Pydantic models for LLM provider configuration, model management,
and chat completions for the internal LLM Chat feature.
The schemas support:
- LLM provider CRUD operations
- Model configuration and capabilities
- Chat completion requests/responses (OpenAI-compatible)
- Embedding requests/responses
"""
# Standard
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Union
# Third-Party
from pydantic import BaseModel, ConfigDict, Field
# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------
class LLMProviderTypeEnum(str, Enum):
"""Enumeration of supported LLM provider types."""
OPENAI = "openai"
AZURE_OPENAI = "azure_openai"
ANTHROPIC = "anthropic"
BEDROCK = "bedrock"
GOOGLE_VERTEX = "google_vertex"
WATSONX = "watsonx"
OLLAMA = "ollama"
OPENAI_COMPATIBLE = "openai_compatible"
COHERE = "cohere"
MISTRAL = "mistral"
GROQ = "groq"
TOGETHER = "together"
class HealthStatus(str, Enum):
"""Health status values for LLM providers."""
HEALTHY = "healthy"
UNHEALTHY = "unhealthy"
UNKNOWN = "unknown"
class RequestStatus(str, Enum):
"""Request processing status."""
PENDING = "pending"
PROCESSING = "processing"
COMPLETED = "completed"
FAILED = "failed"
class RequestType(str, Enum):
"""Types of LLM requests."""
CHAT = "chat"
COMPLETION = "completion"
EMBEDDING = "embedding"
# ---------------------------------------------------------------------------
# LLM Provider Schemas
# ---------------------------------------------------------------------------
class LLMProviderBase(BaseModel):
"""Base schema for LLM provider data."""
name: str = Field(..., min_length=1, max_length=255, description="Display name for the provider")
description: Optional[str] = Field(None, max_length=2000, description="Optional description")
provider_type: LLMProviderTypeEnum = Field(..., description="Type of LLM provider")
api_base: Optional[str] = Field(None, max_length=512, description="Base URL for API requests")
api_version: Optional[str] = Field(None, max_length=50, description="API version (for Azure OpenAI)")
config: Dict[str, Any] = Field(default_factory=dict, description="Provider-specific configuration")
default_model: Optional[str] = Field(None, max_length=255, description="Default model ID")
default_temperature: float = Field(default=0.7, ge=0.0, le=2.0, description="Default temperature")
default_max_tokens: Optional[int] = Field(None, ge=1, description="Default max tokens")
enabled: bool = Field(default=True, description="Whether provider is enabled")
plugin_ids: List[str] = Field(default_factory=list, description="Attached plugin IDs")
def validate_provider_config(self) -> None:
"""Validate provider-specific configuration based on provider type.
Raises:
ValueError: If required provider-specific fields are missing.
"""
# Import here to avoid circular dependency
# First-Party
from mcpgateway.llm_provider_configs import get_provider_config # pylint: disable=import-outside-toplevel
provider_def = get_provider_config(self.provider_type.value if isinstance(self.provider_type, LLMProviderTypeEnum) else self.provider_type)
if not provider_def:
return
# Validate required config fields
for field_def in provider_def.config_fields:
if field_def.required and field_def.name not in self.config:
raise ValueError(f"Required configuration field '{field_def.name}' missing for provider type '{self.provider_type}'")
class LLMProviderCreate(LLMProviderBase):
"""Schema for creating a new LLM provider."""
api_key: Optional[str] = Field(None, description="API key (will be encrypted)")
class LLMProviderUpdate(BaseModel):
"""Schema for updating an LLM provider."""
name: Optional[str] = Field(None, min_length=1, max_length=255)
description: Optional[str] = Field(None, max_length=2000)
provider_type: Optional[LLMProviderTypeEnum] = None
api_key: Optional[str] = Field(None, description="API key (will be encrypted)")
api_base: Optional[str] = Field(None, max_length=512)
api_version: Optional[str] = Field(None, max_length=50)
config: Optional[Dict[str, Any]] = None
default_model: Optional[str] = Field(None, max_length=255)
default_temperature: Optional[float] = Field(None, ge=0.0, le=2.0)
default_max_tokens: Optional[int] = Field(None, ge=1)
enabled: Optional[bool] = None
plugin_ids: Optional[List[str]] = None
class LLMProviderResponse(BaseModel):
"""Schema for LLM provider response."""
model_config = ConfigDict(from_attributes=True)
id: str
name: str
slug: str
description: Optional[str] = None
provider_type: str
api_base: Optional[str] = None
api_version: Optional[str] = None
config: Dict[str, Any] = Field(default_factory=dict)
default_model: Optional[str] = None
default_temperature: float = 0.7
default_max_tokens: Optional[int] = None
enabled: bool = True
health_status: str = "unknown"
last_health_check: Optional[datetime] = None
plugin_ids: List[str] = Field(default_factory=list)
created_at: datetime
updated_at: datetime
created_by: Optional[str] = None
modified_by: Optional[str] = None
model_count: int = Field(default=0, description="Number of models for this provider")
class LLMProviderListResponse(BaseModel):
"""Schema for paginated list of LLM providers."""
providers: List[LLMProviderResponse]
total: int
page: int = 1
page_size: int = 50
# ---------------------------------------------------------------------------
# LLM Model Schemas
# ---------------------------------------------------------------------------
class LLMModelBase(BaseModel):
"""Base schema for LLM model data."""
model_id: str = Field(..., min_length=1, max_length=255, description="Provider's model ID")
model_name: str = Field(..., min_length=1, max_length=255, description="Display name")
model_alias: Optional[str] = Field(None, max_length=255, description="Optional routing alias")
description: Optional[str] = Field(None, max_length=2000, description="Model description")
supports_chat: bool = Field(default=True, description="Supports chat completions")
supports_streaming: bool = Field(default=True, description="Supports streaming")
supports_function_calling: bool = Field(default=False, description="Supports function/tool calling")
supports_vision: bool = Field(default=False, description="Supports vision/images")
context_window: Optional[int] = Field(None, ge=1, description="Max context tokens")
max_output_tokens: Optional[int] = Field(None, ge=1, description="Max output tokens")
enabled: bool = Field(default=True, description="Whether model is enabled")
deprecated: bool = Field(default=False, description="Whether model is deprecated")
class LLMModelCreate(LLMModelBase):
"""Schema for creating a new LLM model."""
provider_id: str = Field(..., description="Provider ID this model belongs to")
class LLMModelUpdate(BaseModel):
"""Schema for updating an LLM model."""
model_id: Optional[str] = Field(None, min_length=1, max_length=255)
model_name: Optional[str] = Field(None, min_length=1, max_length=255)
model_alias: Optional[str] = Field(None, max_length=255)
description: Optional[str] = Field(None, max_length=2000)
supports_chat: Optional[bool] = None
supports_streaming: Optional[bool] = None
supports_function_calling: Optional[bool] = None
supports_vision: Optional[bool] = None
context_window: Optional[int] = Field(None, ge=1)
max_output_tokens: Optional[int] = Field(None, ge=1)
enabled: Optional[bool] = None
deprecated: Optional[bool] = None
class LLMModelResponse(BaseModel):
"""Schema for LLM model response."""
model_config = ConfigDict(from_attributes=True)
id: str
provider_id: str
model_id: str
model_name: str
model_alias: Optional[str] = None
description: Optional[str] = None
supports_chat: bool = True
supports_streaming: bool = True
supports_function_calling: bool = False
supports_vision: bool = False
context_window: Optional[int] = None
max_output_tokens: Optional[int] = None
enabled: bool = True
deprecated: bool = False
created_at: datetime
updated_at: datetime
provider_name: Optional[str] = Field(None, description="Provider name for display")
provider_type: Optional[str] = Field(None, description="Provider type for display")
class LLMModelListResponse(BaseModel):
"""Schema for paginated list of LLM models."""
models: List[LLMModelResponse]
total: int
page: int = 1
page_size: int = 50
# ---------------------------------------------------------------------------
# Chat Completion Schemas (OpenAI-compatible)
# ---------------------------------------------------------------------------
class FunctionDefinition(BaseModel):
"""Function definition for tool calling."""
name: str = Field(..., description="Function name")
description: Optional[str] = Field(None, description="Function description")
parameters: Dict[str, Any] = Field(default_factory=dict, description="JSON Schema for parameters")
class ToolDefinition(BaseModel):
"""Tool definition for function calling."""
type: Literal["function"] = "function"
function: FunctionDefinition
class ChatMessage(BaseModel):
"""A single chat message."""
role: Literal["system", "user", "assistant", "tool"] = Field(..., description="Message role")
content: Optional[str] = Field(None, description="Message content")
name: Optional[str] = Field(None, description="Optional name for the participant")
tool_calls: Optional[List[Dict[str, Any]]] = Field(None, description="Tool calls made by assistant")
tool_call_id: Optional[str] = Field(None, description="ID of tool call this message responds to")
class ChatCompletionRequest(BaseModel):
"""Request for chat completions (OpenAI-compatible)."""
model: str = Field(..., description="Model ID to use")
messages: List[ChatMessage] = Field(..., min_length=1, description="Conversation messages")
temperature: Optional[float] = Field(None, ge=0.0, le=2.0, description="Sampling temperature")
max_tokens: Optional[int] = Field(None, ge=1, description="Maximum tokens to generate")
stream: bool = Field(default=False, description="Enable streaming response")
tools: Optional[List[ToolDefinition]] = Field(None, description="Available tools")
tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(None, description="Tool choice preference")
top_p: Optional[float] = Field(None, ge=0.0, le=1.0, description="Nucleus sampling")
frequency_penalty: Optional[float] = Field(None, ge=-2.0, le=2.0, description="Frequency penalty")
presence_penalty: Optional[float] = Field(None, ge=-2.0, le=2.0, description="Presence penalty")
stop: Optional[Union[str, List[str]]] = Field(None, description="Stop sequences")
user: Optional[str] = Field(None, description="User identifier")
class UsageStats(BaseModel):
"""Token usage statistics."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
class ChatChoice(BaseModel):
"""A single chat completion choice."""
index: int = 0
message: ChatMessage
finish_reason: Optional[str] = None
class ChatCompletionResponse(BaseModel):
"""Response from chat completions."""
id: str = Field(..., description="Unique response ID")
object: str = "chat.completion"
created: int = Field(..., description="Unix timestamp")
model: str = Field(..., description="Model used")
choices: List[ChatChoice]
usage: Optional[UsageStats] = None
class ChatCompletionChunk(BaseModel):
"""Streaming chunk for chat completions."""
id: str
object: str = "chat.completion.chunk"
created: int
model: str
choices: List[Dict[str, Any]]
# ---------------------------------------------------------------------------
# Embedding Schemas
# ---------------------------------------------------------------------------
class EmbeddingRequest(BaseModel):
"""Request for embeddings."""
model: str = Field(..., description="Model ID to use")
input: Union[str, List[str]] = Field(..., description="Text to embed")
encoding_format: Optional[Literal["float", "base64"]] = Field(None, description="Encoding format")
user: Optional[str] = Field(None, description="User identifier")
class EmbeddingData(BaseModel):
"""A single embedding result."""
object: str = "embedding"
embedding: List[float]
index: int = 0
class EmbeddingResponse(BaseModel):
"""Response from embeddings."""
object: str = "list"
data: List[EmbeddingData]
model: str
usage: UsageStats
# ---------------------------------------------------------------------------
# Gateway Models Response (for LLM Chat dropdown)
# ---------------------------------------------------------------------------
class GatewayModelInfo(BaseModel):
"""Simplified model info for the LLM Chat dropdown."""
model_config = ConfigDict(from_attributes=True)
id: str = Field(..., description="Unique model ID")
model_id: str = Field(..., description="Provider's model identifier")
model_name: str = Field(..., description="Display name")
provider_id: str = Field(..., description="Provider ID")
provider_name: str = Field(..., description="Provider display name")
provider_type: str = Field(..., description="Provider type")
supports_streaming: bool = True
supports_function_calling: bool = False
supports_vision: bool = False
class GatewayModelsResponse(BaseModel):
"""Response for /llmchat/gateway/models endpoint."""
models: List[GatewayModelInfo]
count: int
# ---------------------------------------------------------------------------
# Health Check Schemas
# ---------------------------------------------------------------------------
class ProviderHealthCheck(BaseModel):
"""Result of a provider health check."""
provider_id: str
provider_name: str
provider_type: str
status: HealthStatus
response_time_ms: Optional[float] = None
error: Optional[str] = None
checked_at: datetime
