EvoScientist/tests/registry_fixtures.py

"""Shared builders for model-registry integration tests.

Centralizes the "active registry + verified models + snapshot" fixture
graph so middleware and runtime tests don't each re-derive the RegistryV4
payloads. Mirrors the fixtures in ``tests/test_snapshots.py``; that module
keeps its own copies to stay self-contained (Task 4 deliverable).
"""

from __future__ import annotations

from EvoScientist.model_registry.adapters import find_adapter_spec
from EvoScientist.model_registry.hashing import configuration_hash
from EvoScientist.model_registry.resolver import ModelRegistryResolver
from EvoScientist.model_registry.schemas import (
    CredentialWrite,
    ModelRef,
    RegistryV4,
)
from EvoScientist.model_registry.snapshots import (
    SnapshotCreateRequest,
    SnapshotService,
)
from EvoScientist.model_registry.store import ModelRuntimeStore

ZHIPU_REF = ModelRef(provider_id="zhipu-glm", model_key="glm-5.2")
OLLAMA_REF = ModelRef(provider_id="local-ollama", model_key="qwen3")
ZHIPU_SECRET = "sk-live-9876abcd"

DEFAULT_MODEL_RUNTIME = {
    "limit_mode": "combined",
    "context_window_tokens": 1048576,
    "max_input_tokens": None,
    "max_output_tokens": 32768,
    "min_effective_input_tokens": 8192,
    "fixed_system_reserve_tokens": 4096,
    "fixed_tools_reserve_tokens": 8192,
    "fixed_attachments_reserve_tokens": 4096,
    "limits_status": "confirmed",
    "limits_source": "provider",
    "temperature": None,
    "top_p": None,
    "reasoning_effort": "auto",
    "declared_capabilities": {
        "tools": True,
        "vision": False,
        "structured_output": True,
    },
}


def model_runtime_payload(**overrides) -> dict:
    payload = dict(DEFAULT_MODEL_RUNTIME)
    payload.update(overrides)
    return payload


def zhipu_provider_payload(**model_overrides) -> dict:
    return {
        "id": "zhipu-glm",
        "name": "Zhipu GLM",
        "adapter": "openai-compatible",
        "base_url": "https://open.bigmodel.cn/api/paas/v4",
        "auth": {"mode": "api_key", "credential_id": "zhipu-primary"},
        "enabled": True,
        "runtime": {
            "timeout_seconds": 120,
            "max_retries": 2,
            "default_temperature": 0.7,
            "default_top_p": 0.95,
            "default_reasoning_effort": "auto",
        },
        "models": [
            {
                "key": "glm-5.2",
                "name": "GLM-5.2",
                "upstream_model_id": "glm-5.2",
                "enabled": True,
                "runtime": model_runtime_payload(**model_overrides),
            }
        ],
    }


def ollama_provider_payload(**model_overrides) -> dict:
    return {
        "id": "local-ollama",
        "name": "Local Ollama",
        "adapter": "ollama",
        "base_url": "http://localhost:11434",
        "auth": {"mode": "none", "credential_id": None},
        "enabled": True,
        "runtime": {"timeout_seconds": 120, "max_retries": 2},
        "models": [
            {
                "key": "qwen3",
                "name": "Qwen3",
                "upstream_model_id": "qwen3",
                "enabled": True,
                "runtime": model_runtime_payload(**model_overrides),
            }
        ],
    }


def registry_payload(*, auxiliary_default: bool = True, **model_overrides) -> dict:
    return {
        "version": 4,
        "revision": 1,
        "state": "bootstrap",
        "defaults": {
            "primary": {"provider_id": "zhipu-glm", "model_key": "glm-5.2"},
            "auxiliary": (
                {"provider_id": "local-ollama", "model_key": "qwen3"}
                if auxiliary_default
                else None
            ),
        },
        "providers": [
            zhipu_provider_payload(**model_overrides),
            ollama_provider_payload(**model_overrides),
        ],
    }


def _verify_model(store, registry, provider_id, model_key):
    provider = registry.find_provider(provider_id)
    model = provider.find_model(model_key)
    spec = find_adapter_spec(provider.adapter, model.upstream_model_id)
    store.record_model_verification(
        provider_id=provider_id,
        model_key=model_key,
        configuration_hash=configuration_hash(provider, model),
        credential_revision=1 if provider.auth.credential_id else 0,
        adapter_spec_revision=spec.spec_revision,
        result="passed",
        verified_capabilities={
            "tools": True,
            "vision": False,
            "structured_output": True,
        },
    )


def activate_store(store, *, auxiliary_default: bool = True, **model_overrides):
    """Persist an active registry with verified models into *store*."""
    registry = store.save_registry(
        expected_revision=1,
        registry=RegistryV4.model_validate(
            registry_payload(auxiliary_default=auxiliary_default, **model_overrides)
        ),
        credential_writes=[
            CredentialWrite(credential_id="zhipu-primary", secret_value=ZHIPU_SECRET)
        ],
    )
    assert registry.state == "active"
    _verify_model(store, registry, "zhipu-glm", "glm-5.2")
    _verify_model(store, registry, "local-ollama", "qwen3")
    return store


def make_active_store(config_dir, *, auxiliary_default: bool = True, **model_overrides):
    """Persist an active registry with verified models into a fresh store."""
    return activate_store(
        ModelRuntimeStore(config_dir=config_dir),
        auxiliary_default=auxiliary_default,
        **model_overrides,
    )


def make_snapshot_request(**overrides) -> SnapshotCreateRequest:
    payload = {
        "run_request_id": "req-1",
        "thread_id": "thread-1",
        "deployment_id": "local",
        "model_selection_revision": 0,
        "primary": None,
        "auxiliary": None,
    }
    payload.update(overrides)
    return SnapshotCreateRequest.model_validate(payload)


def make_snapshot(store, **overrides):
    """Create a snapshot against the store's registry defaults."""
    service = SnapshotService(store, ModelRegistryResolver(store))
    return service.create(make_snapshot_request(**overrides)).snapshot