208 lines
7.5 KiB
Python
208 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Migrate legacy inline LLM credentials into centralized model configs."""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from sqlmodel import Session, select
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT / "backend"))
|
|
|
|
from agenteval.services.model_configs import SecretCipher, SecretKeyError # noqa: E402
|
|
from agenteval.storage.db import ( # noqa: E402
|
|
ModelConfigDB,
|
|
ScenarioDB,
|
|
ScenarioModelBindingDB,
|
|
engine,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LegacyConfig:
|
|
capability: str
|
|
endpoint_url: str
|
|
model_name: str | None
|
|
api_key: str | None
|
|
|
|
@property
|
|
def fingerprint(self) -> str:
|
|
source = "\0".join((self.capability, self.endpoint_url, self.model_name or "", self.api_key or ""))
|
|
return hashlib.sha256(source.encode("utf-8")).hexdigest()[:12]
|
|
|
|
|
|
@dataclass
|
|
class ScenarioMigration:
|
|
scenario: ScenarioDB
|
|
configs: dict[str, LegacyConfig]
|
|
|
|
|
|
def _legacy_from_params(capability: str, params: dict[str, Any]) -> LegacyConfig | None:
|
|
endpoint = str(params.get("api_url") or "").strip()
|
|
if not endpoint:
|
|
return None
|
|
model = str(params.get("model") or "").strip() or None
|
|
api_key = str(params.get("api_key") or "").strip() or None
|
|
return LegacyConfig(capability, endpoint, model, api_key)
|
|
|
|
|
|
def _scan_scenario(scenario: ScenarioDB) -> ScenarioMigration | None:
|
|
by_purpose: dict[str, list[LegacyConfig]] = {}
|
|
if scenario.get_llm_config():
|
|
generator = _legacy_from_params("chat", scenario.get_llm_config() or {})
|
|
if generator:
|
|
by_purpose.setdefault("generator", []).append(generator)
|
|
|
|
for case in scenario.get_cases():
|
|
for rule in case.get("eval_rules", []):
|
|
rule_type = rule.get("type")
|
|
params = rule.get("params") or {}
|
|
if rule_type == "llm_score":
|
|
legacy = _legacy_from_params("chat", params)
|
|
if legacy:
|
|
by_purpose.setdefault("judge", []).append(legacy)
|
|
elif rule_type == "semantic_similarity":
|
|
legacy = _legacy_from_params("embedding", params)
|
|
if legacy:
|
|
by_purpose.setdefault("embedding", []).append(legacy)
|
|
elif rule_type == "safety" and params.get("use_moderation_api"):
|
|
legacy = _legacy_from_params("moderation", params)
|
|
if legacy:
|
|
by_purpose.setdefault("moderation", []).append(legacy)
|
|
|
|
configs: dict[str, LegacyConfig] = {}
|
|
for purpose, candidates in by_purpose.items():
|
|
unique = {item.fingerprint: item for item in candidates}
|
|
if len(unique) > 1:
|
|
raise ValueError(f"场景 {scenario.id} 的 {purpose} 用途存在多个不同旧配置，请先手工合并")
|
|
configs[purpose] = next(iter(unique.values()))
|
|
return ScenarioMigration(scenario, configs) if configs else None
|
|
|
|
|
|
def _strip_legacy_fields(scenario: ScenarioDB, migrated_purposes: set[str]) -> None:
|
|
if "generator" in migrated_purposes:
|
|
scenario.set_llm_config(None)
|
|
rule_purposes = {
|
|
"llm_score": "judge",
|
|
"semantic_similarity": "embedding",
|
|
"safety": "moderation",
|
|
}
|
|
cases = scenario.get_cases()
|
|
for case in cases:
|
|
for rule in case.get("eval_rules", []):
|
|
if rule_purposes.get(rule.get("type")) not in migrated_purposes:
|
|
continue
|
|
params = rule.get("params") or {}
|
|
for field in ("api_url", "api_key", "model"):
|
|
params.pop(field, None)
|
|
rule["params"] = params
|
|
scenario.set_cases(cases)
|
|
|
|
|
|
def _config_name(config: LegacyConfig) -> str:
|
|
model = (config.model_name or "endpoint").replace("/", "-")[:40]
|
|
return f"迁移-{config.capability}-{model}-{config.fingerprint[:8]}"
|
|
|
|
|
|
def migrate(session: Session, apply: bool) -> dict[str, int]:
|
|
existing_bindings = {
|
|
(binding.scenario_id, binding.purpose): binding
|
|
for binding in session.exec(select(ScenarioModelBindingDB)).all()
|
|
}
|
|
migrations = [
|
|
migration
|
|
for scenario in session.exec(select(ScenarioDB)).all()
|
|
if (migration := _scan_scenario(scenario)) is not None
|
|
]
|
|
pending = [
|
|
(migration, purpose, config)
|
|
for migration in migrations
|
|
for purpose, config in migration.configs.items()
|
|
if (migration.scenario.id or "", purpose) not in existing_bindings
|
|
]
|
|
summary = {
|
|
"scenarios": len(migrations),
|
|
"bindings": len(pending),
|
|
"configs": len({config.fingerprint for _, _, config in pending}),
|
|
}
|
|
if not apply:
|
|
return summary
|
|
|
|
cipher = SecretCipher()
|
|
configs_by_fingerprint: dict[str, ModelConfigDB] = {}
|
|
existing_by_name = {
|
|
config.name: config for config in session.exec(select(ModelConfigDB)).all()
|
|
}
|
|
try:
|
|
for _, _, legacy in pending:
|
|
if legacy.fingerprint in configs_by_fingerprint:
|
|
continue
|
|
name = _config_name(legacy)
|
|
config = existing_by_name.get(name)
|
|
if config is None:
|
|
config = ModelConfigDB(
|
|
name=name,
|
|
provider="openai_compatible",
|
|
capability=legacy.capability,
|
|
endpoint_url=legacy.endpoint_url,
|
|
model_name=legacy.model_name,
|
|
api_key_encrypted=cipher.encrypt(legacy.api_key),
|
|
enabled=True,
|
|
is_default=False,
|
|
description="由 v0.3.0 旧配置迁移工具创建",
|
|
)
|
|
session.add(config)
|
|
session.flush()
|
|
configs_by_fingerprint[legacy.fingerprint] = config
|
|
|
|
for migration in migrations:
|
|
migrated_purposes: set[str] = set()
|
|
scenario_id = migration.scenario.id or ""
|
|
for purpose, legacy in migration.configs.items():
|
|
binding = existing_bindings.get((scenario_id, purpose))
|
|
if binding is None:
|
|
config = configs_by_fingerprint[legacy.fingerprint]
|
|
session.add(
|
|
ScenarioModelBindingDB(
|
|
scenario_id=scenario_id,
|
|
purpose=purpose,
|
|
model_config_id=config.id or "",
|
|
)
|
|
)
|
|
migrated_purposes.add(purpose)
|
|
_strip_legacy_fields(migration.scenario, migrated_purposes)
|
|
session.add(migration.scenario)
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
return summary
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="迁移场景中的旧模型配置；默认仅预览")
|
|
parser.add_argument("--apply", action="store_true", help="执行迁移写入；不传时仅预览")
|
|
args = parser.parse_args()
|
|
try:
|
|
with Session(engine) as session:
|
|
summary = migrate(session, apply=args.apply)
|
|
except (ValueError, SecretKeyError) as exc:
|
|
print(f"迁移终止: {exc}", file=sys.stderr)
|
|
return 1
|
|
mode = "已完成" if args.apply else "预览"
|
|
print(
|
|
f"{mode}: 场景 {summary['scenarios']} 个，新增绑定 {summary['bindings']} 个，"
|
|
f"模型配置 {summary['configs']} 个"
|
|
)
|
|
if not args.apply:
|
|
print("未写入数据库；确认后使用 --apply 执行")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|