109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
"""AI provider 运行时治理：endpoint 覆盖（DB 优先于 .env/yaml）+ 热重载。
|
||
|
||
override 写 DB → 读 DB 全量刷 spec 同步缓存 → ProviderManager 下次构建读新缓存。
|
||
api_key 明文存 DB（本地单用户边界），响应层脱敏不回值（见 schemas）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
from app.ai_provider.config.loader import load_models_config
|
||
from app.core.errors import AppError
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.models.governance import ModelEndpointOverride
|
||
from app.repositories import model_endpoints as overrides_repo
|
||
from app.schemas.preferences import UpdateEndpointOverrideRequest
|
||
from app.services.audit import record_audit_event
|
||
|
||
# —— endpoint 运行时覆盖（DB 优先于 .env/yaml）——
|
||
# reload_provider_manager 时（异步 service）读 DB 全量刷此同步缓存，
|
||
# load_models_config 仍可独立使用。None 字段表示不覆盖、回退 yaml/env。
|
||
_endpoint_overrides_cache: dict[tuple[str, str], dict[str, Any]] = {}
|
||
|
||
|
||
def set_endpoint_overrides_cache(overrides: list[dict[str, Any]]) -> None:
|
||
"""用 DB 全量覆盖列表刷新模块缓存（同步）。"""
|
||
global _endpoint_overrides_cache
|
||
_endpoint_overrides_cache = {
|
||
(o["logical_model_name"], o["endpoint_name"]): o for o in overrides
|
||
}
|
||
|
||
|
||
def get_endpoint_overrides_cache() -> dict[tuple[str, str], dict[str, Any]]:
|
||
"""读取当前覆盖缓存快照。"""
|
||
return dict(_endpoint_overrides_cache)
|
||
|
||
|
||
async def list_endpoint_overrides() -> list[dict[str, Any]]:
|
||
"""读 DB 全量覆盖，序列化为 dict 列表（供刷缓存 + 响应）。"""
|
||
async with AsyncSessionLocal() as session:
|
||
rows = await overrides_repo.list_overrides(session)
|
||
return [_override_to_dict(r) for r in rows]
|
||
|
||
|
||
async def reload_provider_manager() -> None:
|
||
"""读 DB 全量 → 刷 endpoint 覆盖同步缓存。"""
|
||
overrides = await list_endpoint_overrides()
|
||
set_endpoint_overrides_cache(overrides)
|
||
|
||
|
||
async def update_endpoint_override(
|
||
logical_model_name: str,
|
||
endpoint_name: str,
|
||
payload: UpdateEndpointOverrideRequest,
|
||
) -> dict[str, Any]:
|
||
"""upsert 覆盖（payload 字段全量写入，None 清除字段）+ reload + 审计。
|
||
|
||
校验 (model, endpoint) 存在于 yaml（避免覆盖幽灵 endpoint）。
|
||
"""
|
||
if not _endpoint_exists(logical_model_name, endpoint_name):
|
||
raise AppError("not_found", "endpoint 不存在于模型配置中。")
|
||
async with AsyncSessionLocal() as session:
|
||
override = ModelEndpointOverride(
|
||
logical_model_name=logical_model_name,
|
||
endpoint_name=endpoint_name,
|
||
base_url=payload.base_url,
|
||
api_key=payload.api_key,
|
||
status=payload.status,
|
||
weight=payload.weight,
|
||
priority=payload.priority,
|
||
)
|
||
saved = await overrides_repo.upsert_override(session, override)
|
||
await record_audit_event(
|
||
session,
|
||
action="model_endpoint.override_updated",
|
||
actor="user",
|
||
input_summary=f"{logical_model_name}/{endpoint_name}",
|
||
resource_type="model_endpoint_override",
|
||
resource_id=f"{logical_model_name}/{endpoint_name}",
|
||
)
|
||
await session.commit()
|
||
await session.refresh(saved)
|
||
result = _override_to_dict(saved)
|
||
await reload_provider_manager()
|
||
return result
|
||
|
||
|
||
def _endpoint_exists(logical_model_name: str, endpoint_name: str) -> bool:
|
||
"""校验 (model, endpoint) 在 yaml 配置中存在（避免覆盖幽灵 endpoint）。"""
|
||
models = load_models_config()
|
||
for model in models:
|
||
if model.name != logical_model_name:
|
||
continue
|
||
return any(ep.name == endpoint_name for ep in model.endpoints)
|
||
return False
|
||
|
||
|
||
def _override_to_dict(row: ModelEndpointOverride) -> dict[str, Any]:
|
||
return {
|
||
"logical_model_name": row.logical_model_name,
|
||
"endpoint_name": row.endpoint_name,
|
||
"base_url": row.base_url,
|
||
"api_key": row.api_key,
|
||
"status": row.status,
|
||
"weight": row.weight,
|
||
"priority": row.priority,
|
||
"updated_at": row.updated_at,
|
||
}
|