From 6d0e41b2dc1b9b18d86a0e1ffbc32e1b1efc8780 Mon Sep 17 00:00:00 2001 From: octo-patch <266937838+octo-patch@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:06:40 +0000 Subject: [PATCH] okf: add MiniMax provider to the reference agent Route the reference agent's --model through a new providers module so a MiniMax id (MiniMax-M3 / MiniMax-M2.7) runs on MiniMax via LiteLlm while a Gemini id keeps the existing default path. Support the global and China regional endpoints via KC_MINIMAX_REGION and read the key from MINIMAX_API_KEY. Heavy imports stay lazy so the default path is unchanged. --- okf/README.md | 5 + okf/src/reference_agent/agent.py | 5 +- okf/src/reference_agent/bundle/synthesizer.py | 15 +- okf/src/reference_agent/cli.py | 10 +- okf/src/reference_agent/providers.py | 175 ++++++++++++++++++ okf/tests/test_providers.py | 84 +++++++++ 6 files changed, 286 insertions(+), 8 deletions(-) create mode 100644 okf/src/reference_agent/providers.py create mode 100644 okf/tests/test_providers.py diff --git a/okf/README.md b/okf/README.md index 0e88bcf8..5d7a931b 100644 --- a/okf/README.md +++ b/okf/README.md @@ -96,6 +96,11 @@ python3.13 -m venv .venv - Gemini: set `GEMINI_API_KEY` (AI Studio) **or** use Vertex AI by setting `GOOGLE_GENAI_USE_VERTEXAI=true`, `GOOGLE_CLOUD_PROJECT=`, and `GOOGLE_CLOUD_LOCATION=`. +- MiniMax: pass a MiniMax model to `--model` (`MiniMax-M3` or `MiniMax-M2.7`) + and set `MINIMAX_API_KEY`. Requests go to the global endpoint + (`https://api.minimax.io/v1`) by default; set `KC_MINIMAX_REGION=cn_zh` to use + the China endpoint (`https://api.minimaxi.com/v1`). Gemini stays the default + when `--model` names a Gemini id. ## How the reference agent works diff --git a/okf/src/reference_agent/agent.py b/okf/src/reference_agent/agent.py index 1201ae0c..5b7fb1d5 100644 --- a/okf/src/reference_agent/agent.py +++ b/okf/src/reference_agent/agent.py @@ -5,6 +5,7 @@ from google.adk import Agent from google.adk.tools import FunctionTool +from reference_agent.providers import build_model from reference_agent.tools.bundle_tools import read_existing_doc, write_concept_doc from reference_agent.tools.source_tools import ( list_concepts, @@ -27,7 +28,7 @@ def _load_prompt(filename: str) -> str: def build_bq_agent(model: str = DEFAULT_MODEL) -> Agent: return Agent( name="okf_bq_reference_agent", - model=model, + model=build_model(model), instruction=_load_prompt("reference_instruction.md"), tools=[ FunctionTool(list_concepts), @@ -42,7 +43,7 @@ def build_bq_agent(model: str = DEFAULT_MODEL) -> Agent: def build_web_agent(model: str = DEFAULT_MODEL) -> Agent: return Agent( name="okf_web_ingestion_agent", - model=model, + model=build_model(model), instruction=_load_prompt("web_ingestion_instruction.md"), tools=[ FunctionTool(list_concepts), diff --git a/okf/src/reference_agent/bundle/synthesizer.py b/okf/src/reference_agent/bundle/synthesizer.py index d12c6635..07d10ff1 100644 --- a/okf/src/reference_agent/bundle/synthesizer.py +++ b/okf/src/reference_agent/bundle/synthesizer.py @@ -2,6 +2,8 @@ import logging +from reference_agent import providers + log = logging.getLogger(__name__) _PROMPT_TEMPLATE = """\ @@ -37,11 +39,14 @@ def synthesize_description( ) prompt = _PROMPT_TEMPLATE.format(rel_path=rel_path, contents=contents) try: - from google import genai - - client = genai.Client() - response = client.models.generate_content(model=model, contents=prompt) - text = (getattr(response, "text", None) or "").strip() + if providers.is_minimax_model(model): + text = providers.generate_text(model, prompt) + else: + from google import genai + + client = genai.Client() + response = client.models.generate_content(model=model, contents=prompt) + text = (getattr(response, "text", None) or "").strip() if not text: return _fallback(children) return text.splitlines()[0].strip() diff --git a/okf/src/reference_agent/cli.py b/okf/src/reference_agent/cli.py index 5b75a973..eb966062 100644 --- a/okf/src/reference_agent/cli.py +++ b/okf/src/reference_agent/cli.py @@ -8,10 +8,12 @@ from reference_agent.agent import DEFAULT_MODEL from reference_agent.bundle.paths import parse_concept_id +from reference_agent.providers import known_models from reference_agent.runner import ReferenceRunner from reference_agent.sources.bigquery import BigQuerySource _SOURCES = ("bq",) +KNOWN_MINIMAX_MODELS = known_models() def _build_source(name: str, args: argparse.Namespace): @@ -86,7 +88,13 @@ def _parser() -> argparse.ArgumentParser: enrich.add_argument( "--model", default=DEFAULT_MODEL, - help=f"Gemini model id (default: {DEFAULT_MODEL}).", + help=( + f"Model id (default: {DEFAULT_MODEL}). A Gemini id runs on Gemini; " + f"a MiniMax id ({', '.join(KNOWN_MINIMAX_MODELS)}) runs on MiniMax " + "via its chat-completions endpoint. Set MINIMAX_API_KEY and, to " + "target the China endpoint, KC_MINIMAX_REGION=cn_zh " + "(default global_en)." + ), ) enrich.add_argument( "--web-seed", diff --git a/okf/src/reference_agent/providers.py b/okf/src/reference_agent/providers.py new file mode 100644 index 00000000..c3b12ed6 --- /dev/null +++ b/okf/src/reference_agent/providers.py @@ -0,0 +1,175 @@ +"""MiniMax provider wiring for the reference agent. + +The reference agent defaults to Gemini: `model` is a plain string that +`google.adk.Agent` and `google.genai` route to Gemini. This module adds a +second provider, MiniMax, without disturbing that default. When `model` names +a MiniMax model, `build_model()` returns a `LiteLlm` instance pointed at a +MiniMax chat-completions endpoint; for every other name it returns the string +unchanged, so existing Gemini invocations behave exactly as before. + +MiniMax serves the same models from two regional endpoints -- a global endpoint +(`api.minimax.io`) and a China endpoint (`api.minimaxi.com`) -- selected with +the `KC_MINIMAX_REGION` environment variable (default: `global_en`). The API +key is read from `MINIMAX_API_KEY`. Only the LLM backend changes; BigQuery and +the web tools are untouched. + +Heavy dependencies (`google.adk`, `litellm`) are imported lazily inside the +builder functions so importing this module -- and the pure lookup helpers below +-- never requires them. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +PROVIDER_NAME = "MiniMax" + +_API_KEY_ENV = "MINIMAX_API_KEY" +_REGION_ENV = "KC_MINIMAX_REGION" +DEFAULT_REGION = "global_en" + + +@dataclass(frozen=True) +class Pricing: + """USD per million tokens. `None` means the model omits that tier.""" + + input: float + output: float + cache_read: float | None = None + cache_write: float | None = None + + +@dataclass(frozen=True) +class ModelSpec: + model_id: str + context_window: int + pricing: Pricing + input_modalities: tuple[str, ...] + thinking: tuple[str, ...] + + +@dataclass(frozen=True) +class RegionConfig: + region: str + openai_base_url: str + anthropic_base_url: str + docs_root: str + + +_MODELS: dict[str, ModelSpec] = { + "MiniMax-M3": ModelSpec( + model_id="MiniMax-M3", + context_window=1_000_000, + pricing=Pricing(input=0.6, output=2.4, cache_read=0.12, cache_write=None), + input_modalities=("text", "image", "video"), + thinking=("adaptive", "disabled"), + ), + "MiniMax-M2.7": ModelSpec( + model_id="MiniMax-M2.7", + context_window=204_800, + pricing=Pricing(input=0.3, output=1.2, cache_read=0.06, cache_write=0.375), + input_modalities=("text",), + thinking=("always_on",), + ), +} + +_REGIONS: dict[str, RegionConfig] = { + "global_en": RegionConfig( + region="global_en", + openai_base_url="https://api.minimax.io/v1", + anthropic_base_url="https://api.minimax.io/anthropic", + docs_root="https://platform.minimax.io/docs", + ), + "cn_zh": RegionConfig( + region="cn_zh", + openai_base_url="https://api.minimaxi.com/v1", + anthropic_base_url="https://api.minimaxi.com/anthropic", + docs_root="https://platform.minimaxi.com/docs", + ), +} + +# Case-insensitive lookup from a user-supplied name to the canonical model id. +_CANONICAL = {mid.lower(): mid for mid in _MODELS} + + +def known_models() -> tuple[str, ...]: + """Canonical MiniMax model ids this provider serves.""" + return tuple(_MODELS) + + +def known_regions() -> tuple[str, ...]: + return tuple(_REGIONS) + + +def is_minimax_model(model: str) -> bool: + """True when `model` names a MiniMax model (case-insensitive).""" + return bool(model) and model.strip().lower() in _CANONICAL + + +def canonical_model_id(model: str) -> str: + try: + return _CANONICAL[model.strip().lower()] + except (AttributeError, KeyError): + raise ValueError(f"Unknown MiniMax model: {model!r}") + + +def get_model_spec(model: str) -> ModelSpec: + return _MODELS[canonical_model_id(model)] + + +def resolve_region(region: str | None = None) -> RegionConfig: + """Pick a MiniMax regional endpoint. + + `region=None` reads `KC_MINIMAX_REGION` (default `global_en`). An unknown + region name is a hard error rather than a silent fallback. + """ + name = (region or os.environ.get(_REGION_ENV) or DEFAULT_REGION).strip() + try: + return _REGIONS[name] + except KeyError: + raise ValueError( + f"Unknown MiniMax region {name!r}; expected one of " + f"{', '.join(sorted(_REGIONS))}." + ) + + +def litellm_model_name(model: str) -> str: + """LiteLLM provider route for a MiniMax model on its chat endpoint.""" + return f"openai/{canonical_model_id(model)}" + + +def build_model(model: str, *, region: str | None = None): + """Return an ADK model for `model`. + + MiniMax models become a `LiteLlm` instance bound to the selected regional + endpoint; any other name is returned unchanged so Gemini stays the default. + """ + if not is_minimax_model(model): + return model + from google.adk.models.lite_llm import LiteLlm # lazy: only for MiniMax + + endpoint = resolve_region(region) + return LiteLlm( + model=litellm_model_name(model), + api_base=endpoint.openai_base_url, + api_key=os.environ.get(_API_KEY_ENV), + ) + + +def generate_text(model: str, prompt: str, *, region: str | None = None) -> str: + """One-shot text completion for a MiniMax model via its chat-completions endpoint. + + Mirrors the Gemini `genai` path used by the index synthesizer, letting a + MiniMax run produce directory descriptions without a Gemini client. + """ + import litellm # lazy: only for MiniMax + + endpoint = resolve_region(region) + response = litellm.completion( + model=litellm_model_name(model), + messages=[{"role": "user", "content": prompt}], + api_base=endpoint.openai_base_url, + api_key=os.environ.get(_API_KEY_ENV), + ) + return (response.choices[0].message.content or "").strip() diff --git a/okf/tests/test_providers.py b/okf/tests/test_providers.py new file mode 100644 index 00000000..fa102051 --- /dev/null +++ b/okf/tests/test_providers.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import pytest + +from reference_agent import providers + + +def test_known_models_are_the_two_current_ids(): + assert providers.known_models() == ("MiniMax-M3", "MiniMax-M2.7") + + +@pytest.mark.parametrize( + "name", + ["MiniMax-M3", "minimax-m3", " MiniMax-M2.7 ", "MINIMAX-M2.7"], +) +def test_is_minimax_model_matches_case_insensitively(name): + assert providers.is_minimax_model(name) + + +@pytest.mark.parametrize("name", ["gemini-flash-latest", "", "gpt-4o", "minimax"]) +def test_non_minimax_names_are_rejected(name): + assert not providers.is_minimax_model(name) + + +def test_canonical_id_normalizes_case_and_whitespace(): + assert providers.canonical_model_id(" minimax-m3 ") == "MiniMax-M3" + with pytest.raises(ValueError): + providers.canonical_model_id("gemini-flash-latest") + + +def test_m3_spec_matches_registry(): + spec = providers.get_model_spec("MiniMax-M3") + assert spec.context_window == 1_000_000 + assert spec.pricing.input == 0.6 + assert spec.pricing.output == 2.4 + assert spec.pricing.cache_read == 0.12 + assert spec.pricing.cache_write is None + assert spec.input_modalities == ("text", "image", "video") + assert spec.thinking == ("adaptive", "disabled") + + +def test_m27_spec_matches_registry(): + spec = providers.get_model_spec("MiniMax-M2.7") + assert spec.context_window == 204_800 + assert spec.pricing.input == 0.3 + assert spec.pricing.output == 1.2 + assert spec.pricing.cache_read == 0.06 + assert spec.pricing.cache_write == 0.375 + assert spec.input_modalities == ("text",) + assert spec.thinking == ("always_on",) + + +def test_default_region_is_global_endpoint(): + region = providers.resolve_region() + assert region.region == "global_en" + assert region.openai_base_url == "https://api.minimax.io/v1" + assert region.anthropic_base_url == "https://api.minimax.io/anthropic" + assert region.docs_root == "https://platform.minimax.io/docs" + + +def test_china_region_endpoint(): + region = providers.resolve_region("cn_zh") + assert region.openai_base_url == "https://api.minimaxi.com/v1" + assert region.anthropic_base_url == "https://api.minimaxi.com/anthropic" + assert region.docs_root == "https://platform.minimaxi.com/docs" + + +def test_region_env_override(monkeypatch): + monkeypatch.setenv("KC_MINIMAX_REGION", "cn_zh") + assert providers.resolve_region().region == "cn_zh" + + +def test_unknown_region_is_rejected(): + with pytest.raises(ValueError): + providers.resolve_region("eu_de") + + +def test_litellm_model_name_uses_openai_route(): + assert providers.litellm_model_name("minimax-m3") == "openai/MiniMax-M3" + assert providers.litellm_model_name("MiniMax-M2.7") == "openai/MiniMax-M2.7" + + +def test_build_model_passes_through_non_minimax(): + assert providers.build_model("gemini-flash-latest") == "gemini-flash-latest"