Imported from chmielvu/web-search-mcp (src/kindly_web_search_mcp_server/inference/AGENTS.md). Install upstream with npx skills add chmielvu/web-search-mcp --skill inference. Copyright stays with the author.

AGENTS.md - Unified Inference Subsystem

Centralized LLM, embedding, and reranking model execution through a catalog-driven engine with provider-agnostic adapters.

Key Files

Architecture

types.py (ModelSpec, ModelCapability, LLMGeneration)  ← canonical contracts
    ↓
registry.py (define_model, add_provider, get_model, get_provider)
  - ModelDefinition: canonical_id, display_name, capabilities
  - ProviderConfig:  model_id, base_url, api_key_env, timeout
  - ProviderAdapter: execute callable registered by name
    ↓
chain.py (ChainSpec, register_chain, get_chain)  ← ordered "canonical_id@provider" refs
    ↓
engine.py (execute_with_fallback)  ← retry + timeout + telemetry + dispatch
    ↓
adapters/*.py  ← provider-specific I/O (self-register on import)
bridges/*.py  ← Domain-specific integrations
    ↓
router.py (LLMRouter)

Unified Registry (registry.py)

The registry replaces the old model_registry.py + provider_registry.py with a single self-documenting file. Every model is defined once as a canonical entry, then associated with one or more providers that can serve it.

Model Definition

define_model(
    "gpt-oss-120b",
    display_name="GPT OSS 120B",
    description="Primary worker LLM — fast, cheap, OpenAI-compatible.",
    capabilities={ModelCapability.CHAT, ModelCapability.STRUCTURED_OUTPUT},
)
add_provider(
    "gpt-oss-120b",
    "groq",
    as_openai(
        model_id="openai/gpt-oss-120b",
        base_url="https://api.groq.com/openai/v1",
        api_key_env="GROQ_API_KEY",
        cost_per_1m_input=0.15,
        cost_per_1m_output=0.60,
    ),
)

Chain Reference Format

Chains reference models as "canonical_id@provider":

register_chain(
    "worker_llm",
    [
        "gpt-oss-120b@groq",
        "qwen/qwen3.8-27b@groq",
        "qwen/qwen3.6-27b@groq",
        "gpt-oss-120b@vercel",
        "gpt-oss-120b@huggingface",
    ],
)

Provider Config Helpers

Cross-Provider Model ID Normalization

Different providers expose the same model under different names. The registry provides built-in normalization:

• normalize_model_id("openai/gpt-oss-120b") → "gpt-oss-120b" strips known prefixes
• resolve_model_id("gpt-oss-120b", "groq") → "openai/gpt-oss-120b" returns provider-specific string

Live Provider Model Inventory

The 2026-07-25 provider refresh confirmed these active text-output models (llama entries removed 2026-09-08 — dropped from Groq's roster; qwen/qwen3.8-27b added 2026-09-08):

The existing gpt-oss-120b and gpt-oss-20b entries were corrected to the provider-reported GPT OSS names. Groq's Prompt Guard and GPT OSS Safeguard models are intentionally excluded because they are moderation/safety models; speech-output Orpheus models and audio-input Whisper transcription models are also excluded because the generic OpenAI-compatible adapter exposes chat/text generation only.

Multiple API Keys via Qualified Provider Keys

The same model can be served by the same provider with different API keys or timeouts. Use a : suffix on the provider name to create distinct configurations that share one model definition and one adapter:

define_model("gemini-3.1-flash-lite", capabilities={CHAT, GROUNDING, ...})
add_provider("gemini-3.1-flash-lite", "google",        as_google(..., api_key_env="GEMINI_API_KEY"))
add_provider("gemini-3.1-flash-lite", "google:second", as_google(..., api_key_env="SECOND_GEMINI_API_KEY"))
add_provider("gemini-3.1-flash-lite", "google:rankllm", as_google(..., default_timeout=20.0))

Chain references use the qualified key: "gemini-3.1-flash-lite@google:second". The engine resolves "google:second" → adapter "google" automatically.

The worker and classifier chains use the same pattern for provider key failover. @groq:second reads SECOND_GROQ_API_KEY; the primary key entry remains GROQ_API_KEY. Secondary entries are attempted immediately after their corresponding primary and before cross-provider fallbacks.

This eliminates the old pattern of defining the same model N times with different canonical IDs just to vary the API key or timeout.

Gotcha — google.genai SDK Precedence: When both GOOGLE_API_KEY and GEMINI_API_KEY are present in os.environ, the google.genai SDK automatically defaults to GOOGLE_API_KEY. If GOOGLE_API_KEY belongs to a GCP project without the Generative Language API enabled, requests will fail with 403 SERVICE_DISABLED. Ensure either GOOGLE_API_KEY is enabled for Gemini or only GEMINI_API_KEY is exported.

Provider Adapter Aliasing

Multiple provider names can share one adapter implementation:

register_provider_adapter(ProviderAdapter(name="openai", execute=execute_openai, ...))
register_provider_alias("groq", "openai")

Chat OpenRouter uses the shared openai adapter under the openrouter name; OpenRouter's /rerank HTTP adapter is registered separately as openrouter_rerank so importing adapters cannot overwrite the chat alias.

Fallback Pattern

chain.get_chain("name") → ChainSpec(models=[...])
engine.execute_with_fallback(chain, operation, **kwargs) → tries primary → each fallback → raises ChainExhaustedError

• The engine resolves get_provider(spec.provider) to find the adapter.
• Non-retryable errors (422 validation and local configuration errors) abort immediately; retryable errors (auth 401/403, bad request 400, not found 404, rate limit, quota/payment 402, connection, server error, conflict, and timeout) try the next spec.
• Telemetry (LLMUsage, LLMGeneration) captured per attempt by provider adapters.
• Domain bridges may impose a total chain budget in addition to per-provider timeouts; canceled provider tasks must be awaited and drained before failure is returned to an MCP tool.
• Voyage adapter calls voyageai.Client.rerank (truncation=True, max_retries=0) on a worker thread and serializes {index, relevance_score} from results.

Catalog Chains

Qualified keys: @google:second uses SECOND_GEMINI_API_KEY and transparently accepts the existing GEMINI_SECOND_API_KEY spelling; @google:rankllm uses the rankllm timeout; @groq:second uses SECOND_GROQ_API_KEY.

Key Design Decisions

• No LLMEndpoint — ModelSpec is the single canonical model representation.
• No FallbackChainSpec — ChainSpec with ordered model_spec_ids replaces it.
• No build_worker_endpoints — Use build_worker_router() instead.
• LLMGeneration.spec — replaces LLMGeneration.endpoint. Access with gen.spec.provider / gen.spec.model_id.
• Provider adapters self-register — imported by adapters/__init__.py to call register_provider_adapter().
• Models register from settings — catalog._register_all() is called at module import time.
• registry.py uses RLock (reentrant) because _get_model_def() is called from within locked contexts in add_provider() and get_model().

Run Context

engine.bind_run_context(tool_call_id, operation) sets a thread-local context consumed by downstream LLM calls for OTEL span correlation. engine.reset_run_context(token) restores previous state.

Testing

uv run pytest tests/test_inference_subsystem.py tests/test_llm_router.py -v