mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(keys): clear allowed_models when auto_fetch_models is disabled
Add logic to automatically set allowed_models to null when auto_fetch_models is turned off for a key. Also include auto_fetch_models and locked_models fields in config export/import, and trigger model fetching for keys with auto_fetch_models enabled during import.
This commit is contained in:
@@ -248,8 +248,9 @@ class AdminUpdateEndpointKeyAdapter(AdminApiAdapter):
|
||||
db.commit()
|
||||
db.refresh(key)
|
||||
|
||||
# 如果刚刚开启了 auto_fetch_models,立即触发一次模型获取
|
||||
# 处理 auto_fetch_models 的开启和关闭
|
||||
if not auto_fetch_enabled_before and auto_fetch_enabled_after:
|
||||
# 刚刚开启了 auto_fetch_models,立即触发一次模型获取
|
||||
logger.info("[AUTO_FETCH] Key %s 开启自动获取模型,立即触发模型获取", self.key_id)
|
||||
try:
|
||||
from src.services.model.fetch_scheduler import get_model_fetch_scheduler
|
||||
@@ -260,6 +261,12 @@ class AdminUpdateEndpointKeyAdapter(AdminApiAdapter):
|
||||
except Exception as e:
|
||||
logger.error(f"触发模型获取失败: {e}")
|
||||
# 不抛出异常,避免影响 Key 更新操作
|
||||
elif auto_fetch_enabled_before and not auto_fetch_enabled_after:
|
||||
# 关闭了 auto_fetch_models,将 allowed_models 设置为 null(不限制)
|
||||
logger.info("[AUTO_FETCH] Key %s 关闭自动获取模型,清空 allowed_models", self.key_id)
|
||||
key.allowed_models = None
|
||||
db.commit()
|
||||
db.refresh(key)
|
||||
|
||||
# 任何字段更新都清除缓存,确保缓存一致性
|
||||
# 包括 is_active、allowed_models、capabilities 等影响权限和行为的字段
|
||||
|
||||
@@ -764,6 +764,8 @@ class AdminExportConfigAdapter(AdminApiAdapter):
|
||||
"capabilities": key.capabilities,
|
||||
"cache_ttl_minutes": key.cache_ttl_minutes,
|
||||
"max_probe_interval_minutes": key.max_probe_interval_minutes,
|
||||
"auto_fetch_models": key.auto_fetch_models,
|
||||
"locked_models": key.locked_models,
|
||||
"is_active": key.is_active,
|
||||
}
|
||||
)
|
||||
@@ -1123,6 +1125,8 @@ class AdminImportConfigAdapter(AdminApiAdapter):
|
||||
capabilities=key_data.get("capabilities"),
|
||||
cache_ttl_minutes=key_data.get("cache_ttl_minutes", 5),
|
||||
max_probe_interval_minutes=key_data.get("max_probe_interval_minutes", 32),
|
||||
auto_fetch_models=key_data.get("auto_fetch_models", False),
|
||||
locked_models=key_data.get("locked_models"),
|
||||
is_active=key_data.get("is_active", True),
|
||||
health_by_format={},
|
||||
circuit_breaker_by_format={},
|
||||
@@ -1131,6 +1135,12 @@ class AdminImportConfigAdapter(AdminApiAdapter):
|
||||
existing_key_values.add(plaintext_key)
|
||||
stats["keys"]["created"] += 1
|
||||
|
||||
# 如果开启了 auto_fetch_models,记录需要触发获取的 Key ID
|
||||
if key_data.get("auto_fetch_models", False):
|
||||
if "keys_to_fetch" not in stats:
|
||||
stats["keys_to_fetch"] = []
|
||||
stats["keys_to_fetch"].append(new_key.id)
|
||||
|
||||
# 导入 Models
|
||||
for model_data in prov_data.get("models", []):
|
||||
global_model_name = model_data.get("global_model_name")
|
||||
@@ -1238,6 +1248,22 @@ class AdminImportConfigAdapter(AdminApiAdapter):
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.clear_all_caches()
|
||||
|
||||
# 触发开启了 auto_fetch_models 的 Key 的模型获取
|
||||
keys_to_fetch = stats.get("keys_to_fetch", [])
|
||||
if keys_to_fetch:
|
||||
logger.info(f"[AUTO_FETCH] 导入了 {len(keys_to_fetch)} 个开启自动获取模型的 Key,触发模型获取")
|
||||
try:
|
||||
from src.services.model.fetch_scheduler import get_model_fetch_scheduler
|
||||
import asyncio
|
||||
scheduler = get_model_fetch_scheduler()
|
||||
for key_id in keys_to_fetch:
|
||||
asyncio.create_task(scheduler._fetch_models_for_key_by_id(key_id))
|
||||
except Exception as e:
|
||||
logger.error(f"触发模型获取失败: {e}")
|
||||
# 不影响导入成功的返回
|
||||
# 从统计信息中移除内部字段
|
||||
stats.pop("keys_to_fetch", None)
|
||||
|
||||
return {
|
||||
"message": "配置导入成功",
|
||||
"stats": stats,
|
||||
|
||||
Reference in New Issue
Block a user