159 lines
5.2 KiB
Python
159 lines
5.2 KiB
Python
from __future__ import annotations
|
||
|
||
from enum import StrEnum
|
||
from typing import Any
|
||
|
||
from pydantic import ( # noqa: TC002 — Pydantic 模型
|
||
BaseModel,
|
||
Field,
|
||
field_validator,
|
||
model_validator,
|
||
)
|
||
|
||
|
||
class SupportedFormats(StrEnum):
|
||
OPENAI = "openai"
|
||
ANTHROPIC = "anthropic"
|
||
|
||
|
||
_BASE_CAPABILITIES = {
|
||
("openai", "gpt"): {"chat", "tool_use", "structured_output"},
|
||
("openai", "gpt-4"): {"chat", "tool_use", "structured_output", "vision", "reasoning"},
|
||
("anthropic", "claude"): {
|
||
"chat",
|
||
"tool_use",
|
||
"structured_output",
|
||
"vision",
|
||
"reasoning",
|
||
"long_context",
|
||
},
|
||
("deepseek", "deepseek"): {"chat", "tool_use", "structured_output", "reasoning"},
|
||
("deepseek", "deepseek-v4"): {
|
||
"chat",
|
||
"tool_use",
|
||
"structured_output",
|
||
"reasoning",
|
||
"long_context",
|
||
},
|
||
}
|
||
|
||
|
||
def derive_capabilities(provider: str, model: str) -> set[str]:
|
||
"""根据 provider + model 前缀自动推导能力。
|
||
|
||
调用者可以在调用此函数后，在 ModelConfig 上覆盖个别标记。
|
||
"""
|
||
for (prov_prefix, model_prefix), caps in _BASE_CAPABILITIES.items():
|
||
if provider.startswith(prov_prefix) and model.startswith(model_prefix):
|
||
return set(caps)
|
||
return {"chat"}
|
||
|
||
|
||
class ModelEndpointConfig(BaseModel):
|
||
name: str
|
||
model: str
|
||
base_url: str | None = None
|
||
base_url_env: str | None = None
|
||
api_key: str | None = None
|
||
api_key_env: str | None = None
|
||
weight: int = 100
|
||
priority: int = 100
|
||
rpm_limit: int | None = None
|
||
tpm_limit: int | None = None
|
||
max_concurrency: int | None = None
|
||
request_timeout_seconds: int | None = None
|
||
status: str = "active"
|
||
extra_params: dict[str, Any] = Field(default_factory=dict)
|
||
|
||
|
||
class ModelConfig(BaseModel):
|
||
name: str
|
||
display_name: str | None = None
|
||
provider: str
|
||
model: str
|
||
format: SupportedFormats # "openai" | "anthropic"
|
||
chat_format: str | None = None
|
||
tags: list[str] = Field(default_factory=list)
|
||
supports_tools: bool | None = None
|
||
supports_vision: bool | None = None
|
||
supports_reasoning_effort: bool | None = None
|
||
supports_thinking: bool | None = None
|
||
supports_streaming: bool | None = True
|
||
supports_structured_output: bool | None = None
|
||
context_window: int = 32000
|
||
response_reserve: int = 4096
|
||
max_output_tokens: int | None = None
|
||
max_thinking_tokens: int | None = None
|
||
reasoning_effort: str | None = None
|
||
timeout_seconds: int | None = 60
|
||
dimensions: int | None = None # 向量维度（embedding 模型专用）
|
||
routing_strategy: str = "weighted_round_robin"
|
||
endpoint_pool_enabled: bool = False
|
||
weight: int = 100
|
||
endpoints: list[ModelEndpointConfig] = Field(default_factory=list)
|
||
|
||
@model_validator(mode="after")
|
||
def _validate_embedding_dimensions(self) -> ModelConfig:
|
||
"""embedding 模型必须声明 dimensions。"""
|
||
if "embedding" in self.tags and self.dimensions is None:
|
||
raise ValueError(
|
||
f"Embedding 模型 '{self.name}' 必须设置 dimensions（如 dimensions: 1024）"
|
||
)
|
||
return self
|
||
|
||
def resolved_capabilities(self) -> dict[str, bool]:
|
||
base = derive_capabilities(self.provider, self.model)
|
||
return {
|
||
"chat": True,
|
||
"tool_use": self.supports_tools
|
||
if self.supports_tools is not None
|
||
else ("tool_use" in base),
|
||
"vision": self.supports_vision
|
||
if self.supports_vision is not None
|
||
else ("vision" in base),
|
||
"reasoning": self.supports_reasoning_effort
|
||
if self.supports_reasoning_effort is not None
|
||
else ("reasoning" in base),
|
||
"thinking": self.supports_thinking
|
||
if self.supports_thinking is not None
|
||
else ("reasoning" in base),
|
||
"streaming": self.supports_streaming if self.supports_streaming is not None else True,
|
||
"structured_output": self.supports_structured_output
|
||
if self.supports_structured_output is not None
|
||
else ("structured_output" in base),
|
||
}
|
||
|
||
|
||
class ParameterMapping(BaseModel):
|
||
provider_param: str
|
||
min: float | None = None
|
||
max: float | None = None
|
||
default: Any = None
|
||
transform: str | None = None # "multiply_1000", "int", "str"
|
||
alias_for: str | None = None # map from another param name
|
||
handler: str | None = None # named handler for complex param transform
|
||
|
||
|
||
class ProviderConfig(BaseModel):
|
||
name: str
|
||
format: SupportedFormats
|
||
base_url_env: str
|
||
api_key_env: str
|
||
base_url: str | None = None
|
||
timeout_seconds: int = 60
|
||
max_retries: int = 1
|
||
parameter_mappings: dict[str, ParameterMapping] = Field(default_factory=dict)
|
||
parameter_handlers: list[str] = Field(
|
||
default_factory=list
|
||
) # ordered global handlers per provider
|
||
extra_headers: dict[str, str] = Field(default_factory=dict)
|
||
supports_streaming: bool = True
|
||
|
||
@field_validator("format")
|
||
@classmethod
|
||
def _validate_format(cls, v: str) -> str:
|
||
allowed = {e.value for e in SupportedFormats}
|
||
if v not in allowed:
|
||
raise ValueError(f"format must be one of {allowed}, got '{v}'")
|
||
return v
|