438 lines
16 KiB
Python
438 lines
16 KiB
Python
"""LlmConfigService — runtime CRUD for LLM providers/fallbacks (U5).
|
|
|
|
This module wraps the YAML read/write logic from
|
|
:mod:`agentkit.server.routes.settings` as a reusable service. Web UI
|
|
routes (``/api/v1/admin/llm/*``) and the CLI ``agentkit admin llm``
|
|
sub-app both call into :class:`LlmConfigService` rather than touching
|
|
the YAML directly, keeping the validation rules (duplicate-provider,
|
|
fallback-in-use guard, API-key masking) in one place.
|
|
|
|
The service writes back to ``agentkit.yaml`` using the existing
|
|
:func:`_write_yaml_config` helper (which preserves ``${VAR}`` env-var
|
|
references and comments via ruamel.yaml). Concurrent writes are
|
|
serialized with :func:`fcntl.flock` to prevent corruption.
|
|
|
|
The service is a module-level singleton (see
|
|
:func:`get_llm_config_service`) so tests can inject a custom instance
|
|
via :func:`set_llm_config_service`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import fcntl
|
|
import logging
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import yaml
|
|
|
|
from agentkit.server.routes.settings import (
|
|
_mask_api_key,
|
|
_read_yaml_config,
|
|
_write_env_var,
|
|
_write_yaml_config,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
#: Regex matching ``${ENV_VAR}`` references in YAML api_key fields.
|
|
_ENV_REF_RE = re.compile(r"^\$\{([^}]+)\}$")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _provider_to_dict(name: str, pconf: dict[str, Any]) -> dict[str, Any]:
|
|
"""Convert a raw YAML provider dict to a masked response dict.
|
|
|
|
The ``api_key`` field is masked via :func:`_mask_api_key`. The
|
|
``models`` field is preserved as-is (it's already a dict in YAML).
|
|
"""
|
|
raw_key = pconf.get("api_key", "")
|
|
# Resolve ${VAR} references for masking — we want to show the
|
|
# masked *resolved* value, not the literal "${VAR}" string.
|
|
resolved_key = _resolve_env_var(raw_key)
|
|
return {
|
|
"name": name,
|
|
"type": pconf.get("type", "openai"),
|
|
"api_key": _mask_api_key(resolved_key),
|
|
"base_url": pconf.get("base_url", ""),
|
|
"models": pconf.get("models", {}) or {},
|
|
"max_tokens": pconf.get("max_tokens", 4096),
|
|
"timeout": pconf.get("timeout", 60.0),
|
|
}
|
|
|
|
|
|
def _resolve_env_var(value: str) -> str:
|
|
"""Resolve a ``${VAR}`` reference to its env value (if set).
|
|
|
|
If ``value`` is ``${OPENAI_API_KEY}`` and ``OPENAI_API_KEY`` is set
|
|
in the environment, returns the env value. Otherwise returns the
|
|
input unchanged.
|
|
"""
|
|
if not isinstance(value, str):
|
|
return ""
|
|
match = _ENV_REF_RE.match(value)
|
|
if match:
|
|
var_name = match.group(1).split(":-")[0]
|
|
return os.environ.get(var_name, value)
|
|
return value
|
|
|
|
|
|
def _env_var_name_for_provider(name: str) -> str:
|
|
"""Derive a sensible env var name from a provider name."""
|
|
return f"{name.upper().replace('-', '_')}_API_KEY"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# File-locking context manager
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _YamlLock:
|
|
"""Context manager that acquires an exclusive ``fcntl.flock`` on the
|
|
YAML file's sibling lockfile (``<config_path>.lock``).
|
|
|
|
The lockfile is created on demand and is not removed after release
|
|
(it's reused on subsequent writes). Holding the lock prevents
|
|
concurrent writers from corrupting the YAML file.
|
|
"""
|
|
|
|
def __init__(self, config_path: Path) -> None:
|
|
self._lock_path = config_path.with_suffix(config_path.suffix + ".lock")
|
|
self._fh = None
|
|
|
|
def __enter__(self) -> "_YamlLock":
|
|
self._fh = open(self._lock_path, "w")
|
|
fcntl.flock(self._fh.fileno(), fcntl.LOCK_EX)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
if self._fh is not None:
|
|
fcntl.flock(self._fh.fileno(), fcntl.LOCK_UN)
|
|
self._fh.close()
|
|
self._fh = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Service
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class LlmConfigService:
|
|
"""CRUD for LLM providers, API keys, and fallback chains.
|
|
|
|
All methods read/write the YAML file at ``self._config_path``.
|
|
Writes are serialized via :class:`_YamlLock` (``fcntl.flock``) to
|
|
prevent concurrent corruption. The existing ``watch_config()``
|
|
mechanism detects the file change and rebuilds the
|
|
:class:`LLMGateway` — no explicit notification is needed.
|
|
|
|
API keys are NEVER returned in full — every response masks them via
|
|
:func:`_mask_api_key`. New plaintext keys are written to ``.env``
|
|
(next to the YAML) and the YAML stores a ``${ENV_VAR}`` reference.
|
|
"""
|
|
|
|
def __init__(self, config_path: Path) -> None:
|
|
self._config_path = Path(config_path)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Provider CRUD
|
|
# ------------------------------------------------------------------
|
|
|
|
def list_providers(self) -> list[dict[str, Any]]:
|
|
"""Return all providers with masked API keys."""
|
|
data = self._read()
|
|
providers = data.get("llm", {}).get("providers", {}) or {}
|
|
return [_provider_to_dict(name, pconf) for name, pconf in providers.items()]
|
|
|
|
def get_provider(self, name: str) -> dict[str, Any] | None:
|
|
"""Return a single provider by name (masked key), or ``None``."""
|
|
data = self._read()
|
|
pconf = data.get("llm", {}).get("providers", {}).get(name)
|
|
if pconf is None:
|
|
return None
|
|
return _provider_to_dict(name, pconf)
|
|
|
|
def create_provider(
|
|
self,
|
|
name: str,
|
|
provider_type: str,
|
|
api_key: str,
|
|
base_url: str = "",
|
|
models: dict[str, Any] | None = None,
|
|
max_tokens: int = 4096,
|
|
timeout: float = 60.0,
|
|
) -> dict[str, Any]:
|
|
"""Add a new provider to the YAML file.
|
|
|
|
Args:
|
|
name: Provider name (must be unique within the YAML).
|
|
provider_type: Provider type (``openai`` / ``anthropic`` /
|
|
``gemini`` / ``doubao`` / ``wenxin`` / ``yuanbao``).
|
|
api_key: Plaintext API key. Stored in ``.env`` as
|
|
``{NAME}_API_KEY``; the YAML stores ``${...}`` ref.
|
|
base_url: Optional base URL override.
|
|
models: Optional models dict (e.g. ``{"gpt-4o": {}}``).
|
|
max_tokens: Default max tokens for completions.
|
|
timeout: Request timeout in seconds.
|
|
|
|
Returns:
|
|
The newly-created provider dict (masked key).
|
|
|
|
Raises:
|
|
ValueError: If a provider with the same name already exists.
|
|
"""
|
|
with _YamlLock(self._config_path):
|
|
data = self._read()
|
|
data.setdefault("llm", {}).setdefault("providers", {})
|
|
if name in data["llm"]["providers"]:
|
|
raise ValueError(f"Provider {name!r} already exists")
|
|
|
|
env_key = _env_var_name_for_provider(name)
|
|
_write_env_var(str(self._config_path), env_key, api_key)
|
|
|
|
pconf: dict[str, Any] = {
|
|
"type": provider_type,
|
|
"api_key": f"${{{env_key}}}",
|
|
"base_url": base_url,
|
|
"max_tokens": max_tokens,
|
|
"timeout": timeout,
|
|
}
|
|
if models is not None:
|
|
pconf["models"] = models
|
|
data["llm"]["providers"][name] = pconf
|
|
self._write(data)
|
|
|
|
created = self.get_provider(name)
|
|
assert created is not None # we just inserted it
|
|
return created
|
|
|
|
def update_provider(
|
|
self,
|
|
name: str,
|
|
provider_type: str | None = None,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
models: dict[str, Any] | None = None,
|
|
max_tokens: int | None = None,
|
|
timeout: float | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Partially update a provider.
|
|
|
|
Only the provided fields are updated. If ``api_key`` is a masked
|
|
value (starts with ``****``), it is ignored — the existing key
|
|
is preserved.
|
|
|
|
Raises:
|
|
ValueError: If the provider does not exist.
|
|
"""
|
|
with _YamlLock(self._config_path):
|
|
data = self._read()
|
|
providers = data.get("llm", {}).get("providers", {}) or {}
|
|
if name not in providers:
|
|
raise ValueError(f"Provider {name!r} not found")
|
|
|
|
pconf = providers[name]
|
|
if provider_type is not None:
|
|
pconf["type"] = provider_type
|
|
if base_url is not None:
|
|
pconf["base_url"] = base_url
|
|
if models is not None:
|
|
pconf["models"] = models
|
|
if max_tokens is not None:
|
|
pconf["max_tokens"] = max_tokens
|
|
if timeout is not None:
|
|
pconf["timeout"] = timeout
|
|
if api_key is not None and not api_key.startswith("****"):
|
|
# New plaintext key — write to .env, keep ${VAR} ref in YAML.
|
|
existing_key = str(pconf.get("api_key", ""))
|
|
env_match = _ENV_REF_RE.match(existing_key)
|
|
if env_match:
|
|
env_key = env_match.group(1).split(":-")[0]
|
|
else:
|
|
env_key = _env_var_name_for_provider(name)
|
|
_write_env_var(str(self._config_path), env_key, api_key)
|
|
pconf["api_key"] = f"${{{env_key}}}"
|
|
|
|
providers[name] = pconf
|
|
data.setdefault("llm", {})["providers"] = providers
|
|
self._write(data)
|
|
|
|
updated = self.get_provider(name)
|
|
assert updated is not None
|
|
return updated
|
|
|
|
def delete_provider(self, name: str) -> bool:
|
|
"""Remove a provider from the YAML file.
|
|
|
|
Returns:
|
|
``True`` if the provider was deleted.
|
|
|
|
Raises:
|
|
ValueError: If the provider does not exist, or if it is
|
|
referenced by any fallback chain (removing it would
|
|
break the chain).
|
|
"""
|
|
with _YamlLock(self._config_path):
|
|
data = self._read()
|
|
providers = data.get("llm", {}).get("providers", {}) or {}
|
|
if name not in providers:
|
|
raise ValueError(f"Provider {name!r} not found")
|
|
|
|
# Guard: refuse to delete if referenced in any fallback chain.
|
|
fallbacks = data.get("llm", {}).get("fallbacks", {}) or {}
|
|
for model, chain in fallbacks.items():
|
|
if isinstance(chain, list) and name in chain:
|
|
raise ValueError(
|
|
f"Provider {name!r} is used in fallback chain for model "
|
|
f"{model!r}; remove the fallback first"
|
|
)
|
|
|
|
del providers[name]
|
|
data.setdefault("llm", {})["providers"] = providers
|
|
self._write(data)
|
|
return True
|
|
|
|
# ------------------------------------------------------------------
|
|
# API key management
|
|
# ------------------------------------------------------------------
|
|
|
|
def set_api_key(self, provider_name: str, api_key: str) -> dict[str, Any]:
|
|
"""Set the API key for a provider.
|
|
|
|
The plaintext key is written to ``.env`` (as
|
|
``{NAME}_API_KEY``) and the YAML stores a ``${...}`` reference.
|
|
If the provider does not yet exist in the YAML, a stub entry is
|
|
created with ``type=openai`` and default settings.
|
|
|
|
Returns:
|
|
The updated provider dict (masked key).
|
|
"""
|
|
with _YamlLock(self._config_path):
|
|
data = self._read()
|
|
data.setdefault("llm", {}).setdefault("providers", {})
|
|
providers = data["llm"]["providers"]
|
|
pconf = providers.get(provider_name, {})
|
|
existing_key = str(pconf.get("api_key", ""))
|
|
env_match = _ENV_REF_RE.match(existing_key)
|
|
if env_match:
|
|
env_key = env_match.group(1).split(":-")[0]
|
|
else:
|
|
env_key = _env_var_name_for_provider(provider_name)
|
|
_write_env_var(str(self._config_path), env_key, api_key)
|
|
pconf["api_key"] = f"${{{env_key}}}"
|
|
if "type" not in pconf:
|
|
pconf["type"] = "openai"
|
|
providers[provider_name] = pconf
|
|
self._write(data)
|
|
|
|
updated = self.get_provider(provider_name)
|
|
assert updated is not None
|
|
return updated
|
|
|
|
# ------------------------------------------------------------------
|
|
# Fallback chain management
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_fallbacks(self) -> dict[str, list[str]]:
|
|
"""Return the fallback chains (model → list of provider/model refs)."""
|
|
data = self._read()
|
|
fallbacks = data.get("llm", {}).get("fallbacks", {}) or {}
|
|
# Normalize: ensure all values are lists.
|
|
return {
|
|
str(model): list(chain) if isinstance(chain, list) else [str(chain)]
|
|
for model, chain in fallbacks.items()
|
|
}
|
|
|
|
def set_fallback(self, model: str, chain: list[str]) -> dict[str, Any]:
|
|
"""Set the fallback chain for a model.
|
|
|
|
Returns:
|
|
``{"model": model, "chain": chain}`` dict.
|
|
"""
|
|
with _YamlLock(self._config_path):
|
|
data = self._read()
|
|
data.setdefault("llm", {}).setdefault("fallbacks", {})
|
|
data["llm"]["fallbacks"][model] = list(chain)
|
|
self._write(data)
|
|
return {"model": model, "chain": list(chain)}
|
|
|
|
def delete_fallback(self, model: str) -> bool:
|
|
"""Delete the fallback chain for a model.
|
|
|
|
Returns:
|
|
``True`` if a chain was deleted, ``False`` if no chain
|
|
existed for ``model``.
|
|
"""
|
|
with _YamlLock(self._config_path):
|
|
data = self._read()
|
|
fallbacks = data.get("llm", {}).get("fallbacks", {}) or {}
|
|
if model not in fallbacks:
|
|
return False
|
|
del fallbacks[model]
|
|
data.setdefault("llm", {})["fallbacks"] = fallbacks
|
|
self._write(data)
|
|
return True
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internal helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _read(self) -> dict[str, Any]:
|
|
"""Read the YAML config (returns ``{}`` if file is missing/empty)."""
|
|
if not self._config_path.exists():
|
|
return {}
|
|
return _read_yaml_config(str(self._config_path))
|
|
|
|
def _write(self, data: dict[str, Any]) -> None:
|
|
"""Write the full config dict back to the YAML file."""
|
|
_write_yaml_config(str(self._config_path), data)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Module-level singleton (overridable in tests via set_llm_config_service)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_llm_config_service: LlmConfigService | None = None
|
|
|
|
|
|
def get_llm_config_service(config_path: Path | str | None = None) -> LlmConfigService:
|
|
"""Return the process-wide :class:`LlmConfigService`.
|
|
|
|
On first call, ``config_path`` must be provided (or the service
|
|
will raise ``ValueError``). Subsequent calls ignore ``config_path``
|
|
and return the cached singleton.
|
|
"""
|
|
global _llm_config_service
|
|
if _llm_config_service is None:
|
|
if config_path is None:
|
|
raise ValueError("config_path is required to initialize LlmConfigService on first call")
|
|
_llm_config_service = LlmConfigService(Path(config_path))
|
|
return _llm_config_service
|
|
|
|
|
|
def set_llm_config_service(service: LlmConfigService | None) -> None:
|
|
"""Inject a custom :class:`LlmConfigService` (used by tests)."""
|
|
global _llm_config_service
|
|
_llm_config_service = service
|
|
|
|
|
|
# Re-export yaml for tests that need to construct sample config files.
|
|
__all__ = [
|
|
"LlmConfigService",
|
|
"get_llm_config_service",
|
|
"set_llm_config_service",
|
|
"yaml",
|
|
]
|