124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
"""OpenCode Zen model registry for managing model configurations and aliases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.resources
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from utils.env import get_env
|
|
from utils.file_utils import read_json_file
|
|
|
|
from ..shared import ModelCapabilities, ProviderType
|
|
from .base import CAPABILITY_FIELD_NAMES, CapabilityModelRegistry
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ZenModelRegistry(CapabilityModelRegistry):
|
|
"""Capability registry backed by ``conf/zen_models.json``."""
|
|
|
|
LIVE_ENV_VAR_NAME = "ZEN_LIVE_MODELS_CONFIG_PATH"
|
|
LIVE_DEFAULT_FILENAME = "zen_models_live.json"
|
|
|
|
def __init__(self, config_path: str | None = None, live_config_path: str | None = None) -> None:
|
|
self._live_resource = None
|
|
self._live_config_path: Path | None = None
|
|
self._live_default_path = Path(__file__).resolve().parents[3] / "conf" / self.LIVE_DEFAULT_FILENAME
|
|
|
|
if live_config_path:
|
|
self._live_config_path = Path(live_config_path)
|
|
else:
|
|
env_path = get_env(self.LIVE_ENV_VAR_NAME)
|
|
if env_path:
|
|
self._live_config_path = Path(env_path)
|
|
else:
|
|
try:
|
|
resource = importlib.resources.files("conf").joinpath(self.LIVE_DEFAULT_FILENAME)
|
|
if hasattr(resource, "read_text"):
|
|
self._live_resource = resource
|
|
else:
|
|
raise AttributeError("resource accessor not available")
|
|
except Exception:
|
|
self._live_config_path = self._live_default_path
|
|
|
|
super().__init__(
|
|
env_var_name="ZEN_MODELS_CONFIG_PATH",
|
|
default_filename="zen_models.json",
|
|
provider=ProviderType.ZEN,
|
|
friendly_prefix="OpenCode Zen ({model})",
|
|
config_path=config_path,
|
|
)
|
|
|
|
def reload(self) -> None:
|
|
live_data = self._load_live_config_data()
|
|
curated_data = self._load_config_data()
|
|
merged_data = self._merge_manifest_data(live_data, curated_data)
|
|
|
|
self._extras = {}
|
|
configs = [config for config in self._parse_models(merged_data) if config is not None]
|
|
self._build_maps(configs)
|
|
|
|
def _load_live_config_data(self) -> dict:
|
|
if self._live_resource is not None:
|
|
try:
|
|
if hasattr(self._live_resource, "read_text"):
|
|
config_text = self._live_resource.read_text(encoding="utf-8")
|
|
else:
|
|
with self._live_resource.open("r", encoding="utf-8") as handle:
|
|
config_text = handle.read()
|
|
data = json.loads(config_text)
|
|
except FileNotFoundError:
|
|
logger.debug("Packaged %s not found", self.LIVE_DEFAULT_FILENAME)
|
|
return {"models": []}
|
|
except Exception as exc:
|
|
logger.warning("Failed to read packaged %s: %s", self.LIVE_DEFAULT_FILENAME, exc)
|
|
return {"models": []}
|
|
return data or {"models": []}
|
|
|
|
if not self._live_config_path:
|
|
return {"models": []}
|
|
|
|
if not self._live_config_path.exists():
|
|
logger.debug("Zen live registry config not found at %s", self._live_config_path)
|
|
return {"models": []}
|
|
|
|
data = read_json_file(str(self._live_config_path))
|
|
return data or {"models": []}
|
|
|
|
@staticmethod
|
|
def _merge_manifest_data(live_data: dict, curated_data: dict) -> dict:
|
|
merged_models: dict[str, dict] = {}
|
|
|
|
for source in (live_data, curated_data):
|
|
for raw in source.get("models", []):
|
|
if not isinstance(raw, dict):
|
|
continue
|
|
|
|
model_name = raw.get("model_name")
|
|
if not model_name:
|
|
continue
|
|
|
|
existing = merged_models.get(model_name, {})
|
|
merged_models[model_name] = {**existing, **dict(raw)}
|
|
|
|
return {"models": list(merged_models.values())}
|
|
|
|
def _finalise_entry(self, entry: dict) -> tuple[ModelCapabilities, dict]:
|
|
provider_override = entry.get("provider")
|
|
if isinstance(provider_override, str):
|
|
entry_provider = ProviderType(provider_override.lower())
|
|
elif isinstance(provider_override, ProviderType):
|
|
entry_provider = provider_override
|
|
else:
|
|
entry_provider = ProviderType.ZEN
|
|
|
|
entry.setdefault("friendly_name", f"OpenCode Zen ({entry['model_name']})")
|
|
|
|
filtered = {k: v for k, v in entry.items() if k in CAPABILITY_FIELD_NAMES}
|
|
filtered.setdefault("provider", entry_provider)
|
|
capability = ModelCapabilities(**filtered)
|
|
return capability, {}
|