diff --git a/tests/test_litellm_models.py b/tests/test_litellm_models.py new file mode 100644 index 00000000..667d44ec --- /dev/null +++ b/tests/test_litellm_models.py @@ -0,0 +1,109 @@ +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +MODULE_PATH = ( + Path(__file__).parents[1] + / "wdoc" + / "utils" + / "customs" + / "add_extra_litellm_models_metadata.py" +) +SPEC = importlib.util.spec_from_file_location("extra_litellm_metadata", MODULE_PATH) +assert SPEC is not None and SPEC.loader is not None +extra_litellm_metadata = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(extra_litellm_metadata) + + +def test_registers_minimax_models_with_current_metadata(): + registered_models = {} + fake_litellm = SimpleNamespace( + model_cost={}, + models_by_provider={}, + register_model=registered_models.update, + ) + + extra_litellm_metadata.add_extra_models_metadata(fake_litellm) + + expected = { + "minimax/MiniMax-M3": { + "max_tokens": 1_000_000, + "input_cost_per_token": 0.6 / 1_000_000, + "output_cost_per_token": 2.4 / 1_000_000, + "cache_read_input_token_cost": 0.12 / 1_000_000, + "cache_creation_input_token_cost": None, + "input_modalities": ["text", "image", "video"], + "thinking": ["adaptive", "disabled"], + }, + "minimax/MiniMax-M2.7": { + "max_tokens": 204_800, + "input_cost_per_token": 0.3 / 1_000_000, + "output_cost_per_token": 1.2 / 1_000_000, + "cache_read_input_token_cost": 0.06 / 1_000_000, + "cache_creation_input_token_cost": 0.375 / 1_000_000, + "input_modalities": ["text"], + "thinking": ["always_on"], + }, + } + + for model_id, metadata in expected.items(): + assert model_id in fake_litellm.models_by_provider["minimax"] + registered = registered_models[model_id] + assert registered["litellm_provider"] == "minimax" + assert registered["mode"] == "chat" + for key, value in metadata.items(): + assert registered[key] == value + + +def test_does_not_replace_existing_litellm_metadata(): + existing_metadata = {"source": "litellm"} + registered_models = {} + fake_litellm = SimpleNamespace( + model_cost={"minimax/MiniMax-M3": existing_metadata}, + models_by_provider={"minimax": {"minimax/MiniMax-M3"}}, + register_model=registered_models.update, + ) + + extra_litellm_metadata.add_extra_models_metadata(fake_litellm) + + assert fake_litellm.model_cost["minimax/MiniMax-M3"] is existing_metadata + assert "minimax/MiniMax-M3" not in registered_models + assert "minimax/MiniMax-M2.7" in registered_models + + +def test_registration_failure_only_logs_a_warning(monkeypatch): + warnings = [] + fake_litellm = SimpleNamespace(model_cost={}) + monkeypatch.setattr(extra_litellm_metadata.logger, "warning", warnings.append) + monkeypatch.setattr( + extra_litellm_metadata, + "_add_extra_models_metadata", + lambda litellm: (_ for _ in ()).throw(RuntimeError("registration failed")), + ) + + extra_litellm_metadata.add_extra_models_metadata(fake_litellm) + + assert warnings == [ + "Could not add extra LiteLLM model metadata: registration failed" + ] + + +@pytest.mark.parametrize( + ("region", "protocol", "expected"), + [ + ("global_en", "openai", "https://api.minimax.io/v1"), + ("global_en", "anthropic", "https://api.minimax.io/anthropic"), + ("cn_zh", "openai", "https://api.minimaxi.com/v1"), + ("cn_zh", "anthropic", "https://api.minimaxi.com/anthropic"), + ], +) +def test_minimax_endpoint_recipes(region, protocol, expected): + assert extra_litellm_metadata.get_minimax_api_base(region, protocol) == expected + + +def test_minimax_endpoint_recipe_rejects_unknown_selection(): + with pytest.raises(ValueError, match="Unsupported MiniMax endpoint selection"): + extra_litellm_metadata.get_minimax_api_base("unknown", "openai") diff --git a/wdoc/utils/customs/add_extra_litellm_models_metadata.py b/wdoc/utils/customs/add_extra_litellm_models_metadata.py new file mode 100644 index 00000000..84723bc2 --- /dev/null +++ b/wdoc/utils/customs/add_extra_litellm_models_metadata.py @@ -0,0 +1,91 @@ +"""Add model metadata that is not yet available in LiteLLM.""" + +from typing import Any, Literal + +from loguru import logger + + +MINIMAX_PROVIDER = "minimax" + +# Endpoint sources: +# https://platform.minimax.io/docs +# https://platform.minimaxi.com/docs +MINIMAX_ENDPOINTS = { + "global_en": { + "openai_base_url": "https://api.minimax.io/v1", + "anthropic_base_url": "https://api.minimax.io/anthropic", + }, + "cn_zh": { + "openai_base_url": "https://api.minimaxi.com/v1", + "anthropic_base_url": "https://api.minimaxi.com/anthropic", + }, +} + +# Model metadata sources: +# https://platform.minimax.io/docs/api-reference/api-overview +# https://platform.minimaxi.com/docs/api-reference/api-overview +MINIMAX_MODELS = { + "minimax/MiniMax-M3": { + "litellm_provider": MINIMAX_PROVIDER, + "mode": "chat", + "max_tokens": 1_000_000, + "max_input_tokens": 1_000_000, + "input_cost_per_token": 0.6 / 1_000_000, + "output_cost_per_token": 2.4 / 1_000_000, + "cache_read_input_token_cost": 0.12 / 1_000_000, + "cache_creation_input_token_cost": None, + "input_modalities": ["text", "image", "video"], + "thinking": ["adaptive", "disabled"], + "supports_vision": True, + "supports_reasoning": True, + "supports_adaptive_thinking": True, + }, + "minimax/MiniMax-M2.7": { + "litellm_provider": MINIMAX_PROVIDER, + "mode": "chat", + "max_tokens": 204_800, + "max_input_tokens": 204_800, + "input_cost_per_token": 0.3 / 1_000_000, + "output_cost_per_token": 1.2 / 1_000_000, + "cache_read_input_token_cost": 0.06 / 1_000_000, + "cache_creation_input_token_cost": 0.375 / 1_000_000, + "input_modalities": ["text"], + "thinking": ["always_on"], + "supports_reasoning": True, + }, +} + + +def get_minimax_api_base( + region: Literal["global_en", "cn_zh"], + protocol: Literal["openai", "anthropic"] = "openai", +) -> str: + """Return the configured MiniMax base URL for a region and protocol.""" + try: + return MINIMAX_ENDPOINTS[region][f"{protocol}_base_url"] + except KeyError as err: + raise ValueError( + f"Unsupported MiniMax endpoint selection: {region}/{protocol}" + ) from err + + +def _add_extra_models_metadata(litellm: Any) -> None: + models_to_add = { + model_id: metadata + for model_id, metadata in MINIMAX_MODELS.items() + if model_id not in litellm.model_cost + } + if not models_to_add: + return + + litellm.register_model(models_to_add) + provider_models = litellm.models_by_provider.setdefault(MINIMAX_PROVIDER, set()) + provider_models.update(models_to_add) + + +def add_extra_models_metadata(litellm: Any) -> None: + """Add missing model metadata without preventing wdoc from starting.""" + try: + _add_extra_models_metadata(litellm) + except Exception as err: + logger.warning(f"Could not add extra LiteLLM model metadata: {err}") diff --git a/wdoc/wdoc.py b/wdoc/wdoc.py index 75bd16e2..c3cd32aa 100644 --- a/wdoc/wdoc.py +++ b/wdoc/wdoc.py @@ -32,6 +32,9 @@ ) from wdoc.utils.batch_file_loader import batch_load_doc +from wdoc.utils.customs.add_extra_litellm_models_metadata import ( + add_extra_models_metadata, +) from wdoc.utils.env import env, is_out_piped from wdoc.utils.errors import ( NoDocumentsAfterLLMEvalFiltering, @@ -120,6 +123,8 @@ def __init__( """ import litellm + add_extra_models_metadata(litellm) + if version: print(self.VERSION) return