diff --git a/packages/graphrag-llm/graphrag_llm/model_cost_registry/model_cost_registry.py b/packages/graphrag-llm/graphrag_llm/model_cost_registry/model_cost_registry.py index 09a45bc750..d739036532 100644 --- a/packages/graphrag-llm/graphrag_llm/model_cost_registry/model_cost_registry.py +++ b/packages/graphrag-llm/graphrag_llm/model_cost_registry/model_cost_registry.py @@ -9,11 +9,25 @@ from typing_extensions import Self -class ModelCosts(TypedDict): +class ModelCosts(TypedDict, total=False): """Model costs.""" input_cost_per_token: float output_cost_per_token: float + cache_read_input_token_cost: float + max_input_tokens: int + mode: str + + +ADDITIONAL_MODEL_COSTS: dict[str, ModelCosts] = { + "minimax/MiniMax-M3": { + "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, + "max_input_tokens": 1_000_000, + "mode": "chat", + }, +} class ModelCostRegistry: @@ -30,7 +44,7 @@ def __new__(cls, *args: Any, **kwargs: Any) -> Self: def __init__(self): if not hasattr(self, "_initialized"): - self._model_costs = model_cost + self._model_costs = {**model_cost, **ADDITIONAL_MODEL_COSTS} self._initialized = True def register_model_costs(self, model: str, costs: ModelCosts) -> None: diff --git a/tests/integration/language_model/test_factory.py b/tests/integration/language_model/test_factory.py index 428586bf07..91de878cb4 100644 --- a/tests/integration/language_model/test_factory.py +++ b/tests/integration/language_model/test_factory.py @@ -15,6 +15,7 @@ ) from graphrag_llm.config import ModelConfig from graphrag_llm.embedding import LLMEmbedding, create_embedding, register_embedding +from graphrag_llm.model_cost_registry import model_cost_registry if TYPE_CHECKING: from collections.abc import AsyncIterator, Iterator @@ -70,6 +71,17 @@ def tokenizer(self) -> "Tokenizer": ... assert isinstance(model, CustomChatModel) +def test_minimax_m3_model_costs_registered(): + costs = model_cost_registry.get_model_costs("minimax/MiniMax-M3") + + assert costs is not None + assert costs["input_cost_per_token"] == 0.6 / 1_000_000 + assert costs["output_cost_per_token"] == 2.4 / 1_000_000 + assert costs["cache_read_input_token_cost"] == 0.12 / 1_000_000 + assert costs["max_input_tokens"] == 1_000_000 + assert costs["mode"] == "chat" + + def test_create_custom_embedding_llm(): class CustomEmbeddingModel(LLMEmbedding): def __init__(self, **kwargs): ...