137 lines
5.5 KiB
Python
137 lines
5.5 KiB
Python
"""U15 — API Key 加密迁移辅助函数。
|
||
|
||
把 ``agentkit.yaml`` 中 LLM provider 的 plaintext ``api_key`` 迁移到
|
||
:class:`SecretsStore` 加密存储，并更新配置文件中的 ``api_key_encrypted`` /
|
||
``api_key_source`` 字段。
|
||
|
||
迁移模型（双写/双读窗口）：
|
||
1. 读 ``agentkit.yaml`` → ``LLMConfig.providers``。
|
||
2. 对每个 provider 调用 ``ProviderConfig.migrate_to_secrets(store)``：
|
||
- 加密 plaintext 写入 SecretsStore；
|
||
- 验证读回一致后清空 plaintext，标记 ``api_key_source="secrets_store"``；
|
||
- 部分失败时保留 plaintext，标记 ``api_key_source="dual"`` 待重试。
|
||
3. 把更新后的 ``ProviderConfig`` 写回 YAML（plaintext 清空 + encrypted 列写入）。
|
||
4. 返回每个 provider 的迁移状态。
|
||
|
||
回滚步骤（ documented ）：
|
||
- 把 YAML 中 ``api_key_source`` 改回 ``"plaintext"``；
|
||
- 把 plaintext ``api_key`` 重新填回（从备份或 KMS 重新注入）；
|
||
- ``api_key_encrypted`` 列可保留（不影响 plaintext 读取路径）。
|
||
- 重启服务即可。
|
||
|
||
ponytail: CLI 命令接线（``agentkit llm migrate-keys``）延迟实现 —
|
||
本模块的 ``migrate_api_keys_to_secrets`` 函数已可被 CLI / 运维脚本直接调用。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from pathlib import Path
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def migrate_api_keys_to_secrets(config_path: Path | str) -> dict[str, dict[str, object]]:
|
||
"""把 agentkit.yaml 中的 plaintext API Key 迁移到 SecretsStore。
|
||
|
||
流程：
|
||
1. 加载 YAML 配置（不依赖 ServerConfig，直接用 LLMConfig.from_dict）。
|
||
2. 初始化 SecretsStore（master key 从 env ``AGENTKIT_MASTER_KEY`` 读取）。
|
||
3. 对每个 provider 异步执行 ``migrate_to_secrets``。
|
||
4. 把更新后的 providers 段写回 YAML（保留其它段不变）。
|
||
5. 返回 ``{provider_name: {"status": ..., "source": ...}}`` 状态报告。
|
||
|
||
Args:
|
||
config_path: ``agentkit.yaml`` 路径。
|
||
|
||
Returns:
|
||
每个 provider 的迁移状态字典：
|
||
``{"status": "migrated"|"skipped"|"failed", "source": str, "error"?: str}``。
|
||
"""
|
||
import asyncio
|
||
|
||
import yaml
|
||
|
||
from agentkit.channels.secrets import SecretsStore
|
||
from agentkit.llm.config import LLMConfig
|
||
|
||
config_path = Path(config_path)
|
||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {}
|
||
llm_section = raw.get("llm", {})
|
||
llm_config = LLMConfig.from_dict(llm_section)
|
||
|
||
store = SecretsStore() # master key 从 env 加载
|
||
|
||
async def _run() -> dict[str, dict[str, object]]:
|
||
report: dict[str, dict[str, object]] = {}
|
||
for name, pconf in llm_config.providers.items():
|
||
if pconf.api_key_source == "secrets_store" and not pconf.api_key:
|
||
report[name] = {"status": "skipped", "source": pconf.api_key_source}
|
||
continue
|
||
if not pconf.api_key:
|
||
report[name] = {
|
||
"status": "skipped",
|
||
"source": pconf.api_key_source,
|
||
"error": "no plaintext api_key to migrate",
|
||
}
|
||
continue
|
||
try:
|
||
await pconf.migrate_to_secrets(store)
|
||
report[name] = {
|
||
"status": "migrated" if pconf.api_key_source == "secrets_store" else "partial",
|
||
"source": pconf.api_key_source,
|
||
}
|
||
except Exception as e:
|
||
report[name] = {
|
||
"status": "failed",
|
||
"source": pconf.api_key_source,
|
||
"error": str(e),
|
||
}
|
||
return report
|
||
|
||
report = asyncio.run(_run())
|
||
|
||
# 写回 YAML：更新 llm.providers 段，保留其它段
|
||
providers_out: dict[str, dict[str, object]] = {}
|
||
for name, pconf in llm_config.providers.items():
|
||
entry: dict[str, object] = {
|
||
"type": pconf.type,
|
||
"base_url": pconf.base_url,
|
||
"models": pconf.models,
|
||
"max_tokens": pconf.max_tokens,
|
||
"timeout": pconf.timeout,
|
||
"max_connections": pconf.max_connections,
|
||
"max_keepalive_connections": pconf.max_keepalive_connections,
|
||
"keepalive_expiry": pconf.keepalive_expiry,
|
||
# plaintext 已清空（迁移成功）或保留（迁移失败 / dual）
|
||
"api_key": pconf.api_key,
|
||
"api_key_encrypted": pconf.api_key_encrypted,
|
||
"api_key_source": pconf.api_key_source,
|
||
}
|
||
if pconf.retry is not None:
|
||
entry["retry"] = {
|
||
"max_retries": pconf.retry.max_retries,
|
||
"base_delay": pconf.retry.base_delay,
|
||
"max_delay": pconf.retry.max_delay,
|
||
"exponential_base": pconf.retry.exponential_base,
|
||
}
|
||
if pconf.circuit_breaker is not None:
|
||
entry["circuit_breaker"] = {
|
||
"failure_threshold": pconf.circuit_breaker.failure_threshold,
|
||
"recovery_timeout": pconf.circuit_breaker.recovery_timeout,
|
||
"half_open_max": pconf.circuit_breaker.half_open_max,
|
||
}
|
||
providers_out[name] = entry
|
||
|
||
raw.setdefault("llm", {})["providers"] = providers_out
|
||
config_path.write_text(
|
||
yaml.safe_dump(raw, allow_unicode=True, sort_keys=False),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
return report
|
||
|
||
|
||
# ponytail: CLI wiring deferred — 把本函数接到 ``agentkit llm migrate-keys``
|
||
# Typer 命令时，只需在 cli/ 下新增 thin wrapper 调用本函数并打印 report。
|