Skip to content

Refactor models.py to a provider registry #45

Description

@aallan

Context

With #24 expanding to include Mistral, Grok, DeepSeek, and Gemini, we'll have 6+ LLM providers. Currently vera_bench/models.py has three nearly-identical classes:

  • OpenAIClient (native)
  • MoonshotClient (OpenAI-compatible, custom base URL)
  • AnthropicClient (native SDK, different shape)

Adding Mistral/Grok/DeepSeek naively would mean three more copies of MoonshotClient with only the base URL and env var changing. That's ~120 lines of duplication, and any bug fix (e.g., the lazy-import issue we hit in the tests) would need to be made in 5+ places.

The rule of three

With two copies of a class, duplication is tolerable — the pattern hasn't stabilised yet. With three, the pattern is clear and abstraction pays for itself. We're about to add three or four more OpenAI-compatible providers, so this is the right moment to refactor.

Prior art

The Vera compiler is doing exactly this refactor in aallan/vera#413: replacing their _call_inference_provider() elif chain with a _ProviderConfig dataclass and a _PROVIDERS registry dict. Worth mirroring their approach so both repos stay consistent.

Proposed design

from dataclasses import dataclass

@dataclass(frozen=True)
class ProviderConfig:
    env_key: str        # e.g. "MOONSHOT_API_KEY"
    base_url: str       # OpenAI-compat base URL (e.g. "https://api.moonshot.ai/v1")
    prefix: str         # model prefix (e.g. "moonshot/")

_OPENAI_COMPAT_PROVIDERS: dict[str, ProviderConfig] = {
    "moonshot": ProviderConfig(
        env_key="MOONSHOT_API_KEY",
        base_url="https://api.moonshot.ai/v1",
        prefix="moonshot/",
    ),
    "mistral": ProviderConfig(
        env_key="MISTRAL_API_KEY",
        base_url="https://api.mistral.ai/v1",
        prefix="mistral/",
    ),
    "grok": ProviderConfig(
        env_key="XAI_API_KEY",
        base_url="https://api.x.ai/v1",
        prefix="xai/",
    ),
    "deepseek": ProviderConfig(
        env_key="DEEPSEEK_API_KEY",
        base_url="https://api.deepseek.com/v1",
        prefix="deepseek/",
    ),
}

class OpenAICompatClient:
    """Generic client for any OpenAI-compatible API."""
    def __init__(self, model: str, config: ProviderConfig) -> None:
        ...

create_client() iterates _OPENAI_COMPAT_PROVIDERS checking prefixes, returning an OpenAICompatClient with the matching config. Native OpenAIClient and AnthropicClient stay separate (native SDKs). Gemini gets its own class when added.

Scope

  • Add ProviderConfig dataclass and _OPENAI_COMPAT_PROVIDERS registry
  • Add generic OpenAICompatClient class
  • Delete MoonshotClient (replaced by registry entry)
  • Update create_client() to iterate the registry
  • Update tests in tests/test_models.py to cover the generic client with a parametrised test for each provider
  • Update the error message in create_client() to enumerate all registered prefixes

Relates to

  • #24 — Multi-model benchmark runs (blocked on this refactor for clean provider addition)
  • aallan/vera#413 — parallel refactor on the compiler side

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions