373 lines
15 KiB
Python
373 lines
15 KiB
Python
"""Business rules and secret handling for model configurations."""
|
||
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
from typing import Any
|
||
from urllib.parse import urlparse
|
||
|
||
from cryptography.fernet import Fernet, InvalidToken
|
||
from sqlmodel import Session
|
||
|
||
from agenteval.config import get_settings
|
||
from agenteval.model_protocols import ProtocolAdapterError, get_protocol_adapter
|
||
from agenteval.models import ModelCapability, ModelModality, ModelPurpose
|
||
from agenteval.storage.db import ModelConfigDB, ScenarioDB
|
||
from agenteval.storage.model_config_repository import ModelConfigRepository
|
||
|
||
PURPOSE_CAPABILITIES: dict[str, ModelCapability] = {
|
||
ModelPurpose.GENERATOR.value: ModelCapability.CHAT,
|
||
ModelPurpose.JUDGE.value: ModelCapability.CHAT,
|
||
ModelPurpose.EMBEDDING.value: ModelCapability.EMBEDDING,
|
||
ModelPurpose.MODERATION.value: ModelCapability.MODERATION,
|
||
}
|
||
|
||
|
||
class ModelConfigError(ValueError):
|
||
pass
|
||
|
||
|
||
class ModelConfigNotFoundError(ModelConfigError):
|
||
pass
|
||
|
||
|
||
class ModelConfigInUseError(ModelConfigError):
|
||
def __init__(self, references: list[dict[str, str]]):
|
||
super().__init__("模型配置正在被场景引用")
|
||
self.references = references
|
||
|
||
|
||
class SecretKeyError(ModelConfigError):
|
||
pass
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ModelRuntimeConfig:
|
||
id: str
|
||
name: str
|
||
provider: str
|
||
capability: ModelCapability
|
||
endpoint_url: str
|
||
model_name: str | None
|
||
api_key: str | None
|
||
updated_at: datetime | None
|
||
vendor_name: str = ""
|
||
input_modalities: tuple[str, ...] = (ModelModality.TEXT.value,)
|
||
output_modalities: tuple[str, ...] = (ModelModality.TEXT.value,)
|
||
context_window: int | None = None
|
||
max_output_tokens: int | None = None
|
||
supports_streaming: bool = False
|
||
supports_tool_calling: bool = False
|
||
supports_structured_output: bool = False
|
||
supports_reasoning: bool = False
|
||
region: str = ""
|
||
documentation_url: str | None = None
|
||
|
||
def snapshot(self) -> dict[str, Any]:
|
||
return {
|
||
"id": self.id,
|
||
"name": self.name,
|
||
"provider": self.provider,
|
||
"capability": self.capability.value,
|
||
"endpoint_url": self.endpoint_url,
|
||
"model_name": self.model_name,
|
||
"vendor_name": self.vendor_name,
|
||
"input_modalities": list(self.input_modalities),
|
||
"output_modalities": list(self.output_modalities),
|
||
"context_window": self.context_window,
|
||
"max_output_tokens": self.max_output_tokens,
|
||
"supports_streaming": self.supports_streaming,
|
||
"supports_tool_calling": self.supports_tool_calling,
|
||
"supports_structured_output": self.supports_structured_output,
|
||
"supports_reasoning": self.supports_reasoning,
|
||
"region": self.region,
|
||
"documentation_url": self.documentation_url,
|
||
"updated_at": self.updated_at.isoformat() if self.updated_at else None,
|
||
}
|
||
|
||
|
||
class SecretCipher:
|
||
def __init__(self, key: str | None = None):
|
||
self._key = key if key is not None else get_settings().secret_key
|
||
|
||
def _fernet(self) -> Fernet:
|
||
if not self._key:
|
||
raise SecretKeyError("未配置 AGENTEVAL_SECRET_KEY，无法保存或读取模型 API Key")
|
||
try:
|
||
return Fernet(self._key.encode("ascii"))
|
||
except (ValueError, UnicodeEncodeError) as exc:
|
||
raise SecretKeyError("AGENTEVAL_SECRET_KEY 不是有效的 Fernet Key") from exc
|
||
|
||
def encrypt(self, value: str | None) -> str | None:
|
||
if not value:
|
||
return None
|
||
return self._fernet().encrypt(value.encode("utf-8")).decode("ascii")
|
||
|
||
def decrypt(self, value: str | None) -> str | None:
|
||
if not value:
|
||
return None
|
||
try:
|
||
return self._fernet().decrypt(value.encode("ascii")).decode("utf-8")
|
||
except InvalidToken as exc:
|
||
raise SecretKeyError("模型 API Key 无法解密，请检查 AGENTEVAL_SECRET_KEY") from exc
|
||
|
||
|
||
class ModelConfigService:
|
||
def __init__(self, session: Session, cipher: SecretCipher | None = None):
|
||
self.session = session
|
||
self.repo = ModelConfigRepository(session)
|
||
self.cipher = cipher or SecretCipher()
|
||
|
||
@staticmethod
|
||
def validate_fields(provider: str, capability: str, endpoint_url: str, model_name: str | None) -> None:
|
||
try:
|
||
capability_value = ModelCapability(capability)
|
||
except ValueError as exc:
|
||
raise ModelConfigError(f"不支持的模型能力: {capability}") from exc
|
||
try:
|
||
adapter = get_protocol_adapter(provider)
|
||
except ProtocolAdapterError as exc:
|
||
raise ModelConfigError(str(exc)) from exc
|
||
if capability_value not in adapter.supported_capabilities:
|
||
raise ModelConfigError(f"{provider} 协议不支持 {capability_value.value} 能力")
|
||
parsed = urlparse(endpoint_url)
|
||
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
|
||
raise ModelConfigError("Endpoint 必须是有效的 HTTP 或 HTTPS URL")
|
||
if capability_value in {ModelCapability.CHAT, ModelCapability.EMBEDDING} and not model_name:
|
||
raise ModelConfigError("chat 和 embedding 配置必须填写模型名称")
|
||
|
||
@staticmethod
|
||
def validate_common(name: str, provider: str, enabled: bool, is_default: bool) -> None:
|
||
if not name.strip():
|
||
raise ModelConfigError("模型配置名称不能为空")
|
||
try:
|
||
get_protocol_adapter(provider)
|
||
except ProtocolAdapterError as exc:
|
||
raise ModelConfigError(str(exc)) from exc
|
||
if is_default and not enabled:
|
||
raise ModelConfigError("停用的模型配置不能设为默认")
|
||
|
||
@staticmethod
|
||
def normalize_modalities(modalities: list[str] | None) -> list[str]:
|
||
values = modalities if modalities is not None else [ModelModality.TEXT.value]
|
||
normalized: list[str] = []
|
||
for value in values:
|
||
try:
|
||
modality = ModelModality(value).value
|
||
except ValueError as exc:
|
||
raise ModelConfigError(f"不支持的模型模态: {value}") from exc
|
||
if modality not in normalized:
|
||
normalized.append(modality)
|
||
if not normalized:
|
||
raise ModelConfigError("输入和输出模态至少选择一项")
|
||
return normalized
|
||
|
||
@staticmethod
|
||
def validate_metadata(
|
||
context_window: int | None,
|
||
max_output_tokens: int | None,
|
||
documentation_url: str | None,
|
||
) -> None:
|
||
if context_window is not None and context_window <= 0:
|
||
raise ModelConfigError("上下文窗口必须大于 0")
|
||
if max_output_tokens is not None and max_output_tokens <= 0:
|
||
raise ModelConfigError("最大输出 Token 必须大于 0")
|
||
if documentation_url:
|
||
parsed = urlparse(documentation_url)
|
||
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
|
||
raise ModelConfigError("官方文档地址必须是有效的 HTTP 或 HTTPS URL")
|
||
|
||
def create(
|
||
self,
|
||
*,
|
||
name: str,
|
||
provider: str,
|
||
capability: str,
|
||
endpoint_url: str,
|
||
model_name: str | None,
|
||
api_key: str | None,
|
||
enabled: bool,
|
||
is_default: bool,
|
||
description: str,
|
||
vendor_name: str = "",
|
||
input_modalities: list[str] | None = None,
|
||
output_modalities: list[str] | None = None,
|
||
context_window: int | None = None,
|
||
max_output_tokens: int | None = None,
|
||
supports_streaming: bool = False,
|
||
supports_tool_calling: bool = False,
|
||
supports_structured_output: bool = False,
|
||
supports_reasoning: bool = False,
|
||
region: str = "",
|
||
documentation_url: str | None = None,
|
||
) -> ModelConfigDB:
|
||
name = name.strip()
|
||
input_modalities = self.normalize_modalities(input_modalities)
|
||
output_modalities = self.normalize_modalities(output_modalities)
|
||
self.validate_common(name, provider, enabled, is_default)
|
||
self.validate_fields(provider, capability, endpoint_url.strip(), model_name)
|
||
self.validate_metadata(context_window, max_output_tokens, documentation_url)
|
||
if self.repo.get_by_name(name):
|
||
raise ModelConfigError("模型配置名称已存在")
|
||
config = ModelConfigDB(
|
||
name=name,
|
||
provider=provider,
|
||
capability=capability,
|
||
endpoint_url=endpoint_url.strip(),
|
||
model_name=model_name.strip() if model_name else None,
|
||
vendor_name=vendor_name.strip(),
|
||
context_window=context_window,
|
||
max_output_tokens=max_output_tokens,
|
||
supports_streaming=supports_streaming,
|
||
supports_tool_calling=supports_tool_calling,
|
||
supports_structured_output=supports_structured_output,
|
||
supports_reasoning=supports_reasoning,
|
||
region=region.strip(),
|
||
documentation_url=documentation_url.strip() if documentation_url else None,
|
||
api_key_encrypted=self.cipher.encrypt(api_key),
|
||
enabled=enabled,
|
||
is_default=is_default,
|
||
description=description.strip(),
|
||
)
|
||
config.set_modalities(input_modalities, output_modalities)
|
||
return self.repo.create(config)
|
||
|
||
def update(
|
||
self,
|
||
config_id: str,
|
||
*,
|
||
name: str,
|
||
provider: str,
|
||
capability: str,
|
||
endpoint_url: str,
|
||
model_name: str | None,
|
||
api_key: str | None,
|
||
clear_api_key: bool,
|
||
enabled: bool,
|
||
is_default: bool,
|
||
description: str,
|
||
vendor_name: str = "",
|
||
input_modalities: list[str] | None = None,
|
||
output_modalities: list[str] | None = None,
|
||
context_window: int | None = None,
|
||
max_output_tokens: int | None = None,
|
||
supports_streaming: bool = False,
|
||
supports_tool_calling: bool = False,
|
||
supports_structured_output: bool = False,
|
||
supports_reasoning: bool = False,
|
||
region: str = "",
|
||
documentation_url: str | None = None,
|
||
) -> ModelConfigDB:
|
||
config = self.require(config_id)
|
||
name = name.strip()
|
||
input_modalities = self.normalize_modalities(input_modalities)
|
||
output_modalities = self.normalize_modalities(output_modalities)
|
||
self.validate_common(name, provider, enabled, is_default)
|
||
self.validate_fields(provider, capability, endpoint_url.strip(), model_name)
|
||
self.validate_metadata(context_window, max_output_tokens, documentation_url)
|
||
references = self.repo.list_references(config_id)
|
||
if references and not enabled:
|
||
raise ModelConfigError("模型配置正在被场景引用，不能停用")
|
||
for reference in references:
|
||
expected = PURPOSE_CAPABILITIES.get(reference.purpose)
|
||
if expected and expected.value != capability:
|
||
raise ModelConfigError(
|
||
f"模型配置正用于 {reference.purpose}，能力不能改为 {capability}",
|
||
)
|
||
same_name = self.repo.get_by_name(name)
|
||
if same_name and same_name.id != config_id:
|
||
raise ModelConfigError("模型配置名称已存在")
|
||
config.name = name
|
||
config.provider = provider
|
||
config.capability = capability
|
||
config.endpoint_url = endpoint_url.strip()
|
||
config.model_name = model_name.strip() if model_name else None
|
||
config.vendor_name = vendor_name.strip()
|
||
config.set_modalities(input_modalities, output_modalities)
|
||
config.context_window = context_window
|
||
config.max_output_tokens = max_output_tokens
|
||
config.supports_streaming = supports_streaming
|
||
config.supports_tool_calling = supports_tool_calling
|
||
config.supports_structured_output = supports_structured_output
|
||
config.supports_reasoning = supports_reasoning
|
||
config.region = region.strip()
|
||
config.documentation_url = documentation_url.strip() if documentation_url else None
|
||
config.enabled = enabled
|
||
config.is_default = is_default
|
||
config.description = description.strip()
|
||
if clear_api_key:
|
||
config.api_key_encrypted = None
|
||
elif api_key:
|
||
config.api_key_encrypted = self.cipher.encrypt(api_key)
|
||
return self.repo.update(config)
|
||
|
||
def require(self, config_id: str) -> ModelConfigDB:
|
||
config = self.repo.get(config_id)
|
||
if not config:
|
||
raise ModelConfigNotFoundError("模型配置不存在")
|
||
return config
|
||
|
||
def resolve(
|
||
self,
|
||
config_id: str,
|
||
expected_capability: ModelCapability | None = None,
|
||
) -> ModelRuntimeConfig:
|
||
config = self.require(config_id)
|
||
capability = ModelCapability(config.capability)
|
||
if not config.enabled:
|
||
raise ModelConfigError(f"模型配置“{config.name}”已禁用")
|
||
if expected_capability and capability != expected_capability:
|
||
raise ModelConfigError(
|
||
f"模型配置“{config.name}”能力为 {capability.value}，不能用于 {expected_capability.value}",
|
||
)
|
||
return ModelRuntimeConfig(
|
||
id=config.id or "",
|
||
name=config.name,
|
||
provider=config.provider,
|
||
capability=capability,
|
||
endpoint_url=config.endpoint_url,
|
||
model_name=config.model_name,
|
||
api_key=self.cipher.decrypt(config.api_key_encrypted),
|
||
updated_at=config.updated_at,
|
||
vendor_name=config.vendor_name,
|
||
input_modalities=tuple(config.get_input_modalities()),
|
||
output_modalities=tuple(config.get_output_modalities()),
|
||
context_window=config.context_window,
|
||
max_output_tokens=config.max_output_tokens,
|
||
supports_streaming=config.supports_streaming,
|
||
supports_tool_calling=config.supports_tool_calling,
|
||
supports_structured_output=config.supports_structured_output,
|
||
supports_reasoning=config.supports_reasoning,
|
||
region=config.region,
|
||
documentation_url=config.documentation_url,
|
||
)
|
||
|
||
def delete(self, config_id: str) -> None:
|
||
config = self.require(config_id)
|
||
references = []
|
||
for binding in self.repo.list_references(config_id):
|
||
scenario = self.session.get(ScenarioDB, binding.scenario_id)
|
||
references.append(
|
||
{
|
||
"scenario_id": binding.scenario_id,
|
||
"scenario_name": scenario.name if scenario else binding.scenario_id,
|
||
"purpose": binding.purpose,
|
||
}
|
||
)
|
||
if references:
|
||
raise ModelConfigInUseError(references)
|
||
self.repo.delete(config)
|
||
|
||
def validate_bindings(self, bindings: dict[str, str]) -> None:
|
||
for purpose, config_id in bindings.items():
|
||
expected = PURPOSE_CAPABILITIES.get(purpose)
|
||
if expected is None:
|
||
raise ModelConfigError(f"不支持的模型用途: {purpose}")
|
||
config = self.require(config_id)
|
||
if not config.enabled:
|
||
raise ModelConfigError(f"模型配置“{config.name}”已禁用")
|
||
capability = ModelCapability(config.capability)
|
||
if capability != expected:
|
||
raise ModelConfigError(
|
||
f"模型配置“{config.name}”能力为 {capability.value}，不能用于 {expected.value}",
|
||
)
|