AgentEvalTool/scripts/migrate_model_configs.py

#!/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())