mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-06 09:28:34 +00:00
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
"""Helpers for runtime model preset selection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Mapping
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from nanobot.config.schema import ModelPresetConfig
|
|
from nanobot.providers.base import LLMProvider
|
|
from nanobot.providers.factory import ProviderSnapshot, build_provider_snapshot
|
|
|
|
PresetSnapshotLoader = Callable[[str], ProviderSnapshot]
|
|
PresetCatalogLoader = Callable[[], Mapping[str, ModelPresetConfig]]
|
|
|
|
|
|
def default_selection_signature(
|
|
signature: tuple[object, ...] | None,
|
|
model_preset: str | None = None,
|
|
) -> tuple[object, ...] | None:
|
|
return (model_preset, *signature[:2]) if signature else None
|
|
|
|
|
|
def configured_model_presets(config: Any) -> dict[str, ModelPresetConfig]:
|
|
return {**config.model_presets, "default": config.resolve_default_preset()}
|
|
|
|
|
|
def load_model_preset_catalog(
|
|
config_path: Path | None = None,
|
|
) -> dict[str, ModelPresetConfig]:
|
|
"""Load the current preset catalog from the configured file."""
|
|
from nanobot.config.loader import load_config, resolve_config_env_vars
|
|
|
|
return configured_model_presets(
|
|
resolve_config_env_vars(load_config(config_path)),
|
|
)
|
|
|
|
|
|
def make_preset_snapshot_loader(
|
|
config: Any,
|
|
provider_snapshot_loader: Callable[..., ProviderSnapshot] | None,
|
|
) -> PresetSnapshotLoader:
|
|
if provider_snapshot_loader is not None:
|
|
return lambda name: provider_snapshot_loader(preset_name=name)
|
|
return lambda name: build_provider_snapshot(config, preset_name=name)
|
|
|
|
|
|
def build_static_preset_snapshot(
|
|
provider: LLMProvider,
|
|
name: str,
|
|
preset: ModelPresetConfig,
|
|
) -> ProviderSnapshot:
|
|
return ProviderSnapshot(
|
|
provider=provider,
|
|
model=preset.model,
|
|
context_window_tokens=preset.context_window_tokens,
|
|
signature=("model_preset", name, preset.model_dump_json()),
|
|
generation=preset.to_generation_settings(),
|
|
model_preset=name,
|
|
)
|
|
|
|
|
|
def build_runtime_preset_snapshot(
|
|
*,
|
|
name: str,
|
|
presets: dict[str, ModelPresetConfig],
|
|
provider: LLMProvider,
|
|
loader: PresetSnapshotLoader | None,
|
|
) -> ProviderSnapshot:
|
|
if loader is not None:
|
|
return replace(loader(name), model_preset=name)
|
|
return build_static_preset_snapshot(provider, name, presets[name])
|
|
|
|
|
|
def normalize_preset_name(name: str | None, presets: dict[str, ModelPresetConfig]) -> str:
|
|
if not isinstance(name, str) or not name.strip():
|
|
raise ValueError("model_preset must be a non-empty string")
|
|
name = name.strip()
|
|
if name not in presets:
|
|
raise KeyError(f"model_preset {name!r} not found. Available: {', '.join(presets) or '(none)'}")
|
|
return name
|
|
|