258 lines
10 KiB
Python
258 lines
10 KiB
Python
"""LLM Config - 配置加载"""
|
||
|
||
import json
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, TYPE_CHECKING
|
||
|
||
from agentkit.llm.retry import CircuitBreakerConfig, RetryConfig
|
||
|
||
if TYPE_CHECKING:
|
||
from agentkit.channels.secrets import SecretsStore
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class CacheConfig:
|
||
"""LLM Cache 配置"""
|
||
|
||
enabled: bool = False
|
||
backend: str = "auto" # "auto" | "redis" | "memory"
|
||
redis_url: str = "redis://localhost:6379"
|
||
exact_ttl: int = 3600
|
||
semantic_ttl: int = 86400
|
||
similarity_threshold: float = 0.92
|
||
max_entries: int = 10000
|
||
# Embedding config for semantic cache (Chinese-first: bge-m3 via Xinference)
|
||
embedding_provider: str = "openai" # "openai" | "xinference" | "local"
|
||
embedding_model: str = "bge-m3"
|
||
embedding_base_url: str | None = None
|
||
embedding_api_key: str | None = None
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict) -> "CacheConfig":
|
||
if not data:
|
||
return cls()
|
||
emb = data.get("embedding", {})
|
||
return cls(
|
||
enabled=data.get("enabled", False),
|
||
backend=data.get("backend", "auto"),
|
||
redis_url=data.get("redis_url", "redis://localhost:6379"),
|
||
exact_ttl=data.get("exact_ttl", 3600),
|
||
semantic_ttl=data.get("semantic_ttl", 86400),
|
||
similarity_threshold=data.get("similarity_threshold", 0.92),
|
||
max_entries=data.get("max_entries", 10000),
|
||
embedding_provider=emb.get("provider", "openai"),
|
||
embedding_model=emb.get("model", "bge-m3"),
|
||
embedding_base_url=emb.get("base_url"),
|
||
embedding_api_key=emb.get("api_key"),
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class ProviderConfig:
|
||
"""Provider 配置"""
|
||
|
||
api_key: str
|
||
base_url: str
|
||
models: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||
type: str = "openai" # "openai" | "anthropic" | "gemini"
|
||
max_tokens: int = 4096 # Anthropic: default max_tokens
|
||
timeout: float = 120.0 # Anthropic: request timeout
|
||
max_connections: int = 100 # httpx 连接池最大连接数
|
||
max_keepalive_connections: int = 20 # httpx 连接池最大保活连接数
|
||
keepalive_expiry: float = 30.0 # httpx 保活连接过期时间（秒）
|
||
retry: RetryConfig | None = None
|
||
circuit_breaker: CircuitBreakerConfig | None = None
|
||
# U15 — API Key 加密迁移（plaintext → SecretsStore）。
|
||
# api_key_encrypted: JSON 编码的 SecretEntry（base64 nonce+salt+ciphertext）。
|
||
# 为 None 表示未迁移，仍走 plaintext api_key。
|
||
# api_key_source: 当前 key 来源标记，用于双写/双读迁移窗口。
|
||
# "plaintext" — 仅 plaintext 列；"secrets_store" — 仅加密列；
|
||
# "dual" — 双写窗口中（两列都有值，读时优先加密）。
|
||
api_key_encrypted: str | None = None
|
||
api_key_source: str = "plaintext"
|
||
|
||
def get_api_key(self) -> str:
|
||
"""同步读取 API Key — 返回 plaintext。
|
||
|
||
双读窗口的同步入口：无法 await ``SecretsStore.get_secret``，
|
||
因此加密列需通过异步方法 :meth:`aget_api_key` 读取。
|
||
本方法始终返回 plaintext ``api_key``（迁移期保证可用性）。
|
||
"""
|
||
if self.api_key_encrypted:
|
||
logger.debug("get_api_key: encrypted key set — use aget_api_key for decryption")
|
||
return self.api_key
|
||
|
||
async def aget_api_key(self, secrets_store: "SecretsStore | None" = None) -> str:
|
||
"""异步读取 API Key — 双读窗口优先加密列，失败回退 plaintext。
|
||
|
||
Args:
|
||
secrets_store: 可选的加密存储。为 None 时直接返回 plaintext。
|
||
|
||
Returns:
|
||
解密后的 API Key；加密列解密失败时回退到 plaintext ``api_key``。
|
||
"""
|
||
if self.api_key_encrypted and secrets_store is not None:
|
||
try:
|
||
entry = self._decode_secret_entry(self.api_key_encrypted)
|
||
decrypted = await secrets_store.get_secret(entry.key)
|
||
if decrypted is not None:
|
||
return decrypted
|
||
# store 里没有这个 key（可能已被删除）— 回退 plaintext
|
||
logger.warning(
|
||
f"aget_api_key: encrypted key for provider type={self.type} "
|
||
f"not found in secrets_store — fallback to plaintext"
|
||
)
|
||
except Exception as e:
|
||
# 解密失败（master key 不匹配 / 密文损坏）— 回退 plaintext
|
||
logger.warning(
|
||
f"aget_api_key: decrypt failed for type={self.type}: {e} — fallback to plaintext"
|
||
)
|
||
return self.api_key
|
||
|
||
async def migrate_to_secrets(self, secrets_store: "SecretsStore") -> None:
|
||
"""把 plaintext api_key 迁移到 SecretsStore（幂等）。
|
||
|
||
迁移步骤（双写窗口）：
|
||
1. 若 ``api_key_source == "secrets_store"`` 且 plaintext 已清空 → 已迁移，no-op。
|
||
2. 否则：调用 ``secrets_store.set_secret`` 加密存储 key。
|
||
3. 把返回的 SecretEntry JSON 编码写入 ``api_key_encrypted``。
|
||
4. 标记 ``api_key_source = "secrets_store"``。
|
||
5. 验证：调用 ``get_secret`` 读回对比，成功后清空 plaintext ``api_key=""``。
|
||
|
||
幂等性：重复调用不会重复加密（已迁移时直接返回）。
|
||
部分失败恢复：若 set 成功但验证失败，保留 plaintext 不清空，
|
||
``api_key_encrypted`` 已写入 — 下次重试时由幂等性保证最终一致。
|
||
|
||
Args:
|
||
secrets_store: 用于加密存储的 SecretsStore 实例。
|
||
"""
|
||
# 幂等：已迁移完成（source 标记 + plaintext 已清空）
|
||
if self.api_key_source == "secrets_store" and not self.api_key and self.api_key_encrypted:
|
||
return
|
||
|
||
# 没有 plaintext 可迁移（空 key）— 跳过
|
||
if not self.api_key:
|
||
return
|
||
|
||
secret_key = self._secret_key_for_type()
|
||
# 双写：先写加密列
|
||
entry = await secrets_store.set_secret(secret_key, self.api_key)
|
||
self.api_key_encrypted = self._encode_secret_entry(entry, secret_key)
|
||
|
||
# 验证：读回对比
|
||
try:
|
||
decrypted = await secrets_store.get_secret(secret_key)
|
||
except Exception as e:
|
||
logger.warning(f"migrate_to_secrets: verify read failed for type={self.type}: {e}")
|
||
# 加密列已写但验证失败 — 保留 plaintext，标记 dual 待重试
|
||
self.api_key_source = "dual"
|
||
return
|
||
|
||
if decrypted != self.api_key:
|
||
logger.error(
|
||
f"migrate_to_secrets: verify mismatch for type={self.type} "
|
||
f"— plaintext retained, source=dual"
|
||
)
|
||
self.api_key_source = "dual"
|
||
return
|
||
|
||
# 验证通过：清空 plaintext，标记完成
|
||
self.api_key_source = "secrets_store"
|
||
self.api_key = ""
|
||
|
||
def _secret_key_for_type(self) -> str:
|
||
"""生成 SecretsStore 中的 key（按 provider type 命名空间隔离）。"""
|
||
return f"llm:provider:{self.type}:api_key"
|
||
|
||
@staticmethod
|
||
def _encode_secret_entry(entry: Any, key: str) -> str:
|
||
"""把 SecretEntry 编码为 JSON 字符串（含 key 字段）。"""
|
||
# entry 是 SecretEntry pydantic 模型，有 model_dump()
|
||
if hasattr(entry, "model_dump"):
|
||
data = entry.model_dump()
|
||
else:
|
||
data = dict(entry)
|
||
data["key"] = key
|
||
return json.dumps(data)
|
||
|
||
@staticmethod
|
||
def _decode_secret_entry(encoded: str) -> Any:
|
||
"""从 JSON 字符串解码 SecretEntry。返回带 .key 属性的对象。"""
|
||
from agentkit.channels.secrets import SecretEntry
|
||
|
||
data = json.loads(encoded)
|
||
return SecretEntry(
|
||
key=data.get("key", ""),
|
||
value=data["value"],
|
||
nonce=data["nonce"],
|
||
salt=data["salt"],
|
||
key_id=data.get("key_id", "default"),
|
||
created_at=data.get("created_at", ""),
|
||
updated_at=data.get("updated_at", ""),
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class LLMConfig:
|
||
"""LLM 配置"""
|
||
|
||
providers: dict[str, ProviderConfig] = field(default_factory=dict)
|
||
model_aliases: dict[str, str] = field(default_factory=dict)
|
||
fallbacks: dict[str, list[str]] = field(default_factory=dict)
|
||
cache: CacheConfig | None = None
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict) -> "LLMConfig":
|
||
"""从字典加载配置"""
|
||
providers = {}
|
||
for name, pconf in data.get("providers", {}).items():
|
||
retry = None
|
||
retry_data = pconf.get("retry")
|
||
if retry_data:
|
||
retry = RetryConfig(
|
||
max_retries=retry_data.get("max_retries", 3),
|
||
base_delay=retry_data.get("base_delay", 1.0),
|
||
max_delay=retry_data.get("max_delay", 30.0),
|
||
exponential_base=retry_data.get("exponential_base", 2.0),
|
||
)
|
||
|
||
circuit_breaker = None
|
||
cb_data = pconf.get("circuit_breaker")
|
||
if cb_data:
|
||
circuit_breaker = CircuitBreakerConfig(
|
||
failure_threshold=cb_data.get("failure_threshold", 5),
|
||
recovery_timeout=cb_data.get("recovery_timeout", 60.0),
|
||
half_open_max=cb_data.get("half_open_max", 1),
|
||
)
|
||
|
||
providers[name] = ProviderConfig(
|
||
api_key=pconf.get("api_key", ""),
|
||
base_url=pconf.get("base_url", ""),
|
||
models=pconf.get("models", {}),
|
||
type=pconf.get("type", "openai"),
|
||
max_tokens=pconf.get("max_tokens", 4096),
|
||
timeout=pconf.get("timeout", 120.0),
|
||
max_connections=pconf.get("max_connections", 100),
|
||
max_keepalive_connections=pconf.get("max_keepalive_connections", 20),
|
||
keepalive_expiry=pconf.get("keepalive_expiry", 30.0),
|
||
retry=retry,
|
||
circuit_breaker=circuit_breaker,
|
||
# U15 — 新增加密迁移字段，缺省时保持 plaintext 行为
|
||
api_key_encrypted=pconf.get("api_key_encrypted"),
|
||
api_key_source=pconf.get("api_key_source", "plaintext"),
|
||
)
|
||
cache = None
|
||
cache_data = data.get("cache")
|
||
if cache_data:
|
||
cache = CacheConfig.from_dict(cache_data)
|
||
|
||
return cls(
|
||
providers=providers,
|
||
model_aliases=data.get("model_aliases", {}),
|
||
fallbacks=data.get("fallbacks", {}),
|
||
cache=cache,
|
||
)
|