feat: implement get_model TODO and fix critical telemetry bug (#647)

ErikBjare · ellipsis-dev[bot] · web-flow · commit 9c05ac26966a · 2025-09-22T11:26:43.000+02:00
* feat: implement get_model TODO and fix critical telemetry bug

- Enhanced get_model() to use _get_models_for_provider for dynamic model discovery
  - Integrated with existing dynamic fetching infrastructure
  - Added proper fallback handling for unknown providers/models
  - Improved parsing logic to handle provider/model formats correctly

- Fixed critical telemetry bug in OpenAI LLM where stripped model names caused warnings
  - Changed _record_usage calls to pass full model name instead of base_model
  - Resolves "Unknown model x-ai/grok-4-fast:free" warnings during evals
  - Makes OpenAI implementation consistent with Anthropic

- Improved reasoning model detection using metadata instead of hardcoded checks
  - Replaced _is_reasoner() function with model_meta.supports_reasoning
  - Updated extra_body() to use ModelMeta parameter for better reasoning support
  - Enhanced message preparation logic for reasoning models

- Added comprehensive test suite with 10 focused tests
  - Tests static/dynamic model lookup, provider-only requests
  - Validates error handling and fallback scenarios
  - Covers new dynamic model fetching integration

All changes maintain backwards compatibility while significantly enhancing
model discovery capabilities and fixing evaluation warnings.

* Update gptme/llm/models.py

Co-authored-by: ellipsis-dev[bot] &lt;65095814+ellipsis-dev[bot]@users.noreply.github.com&gt;

---------

Co-authored-by: ellipsis-dev[bot] &lt;65095814+ellipsis-dev[bot]@users.noreply.github.com&gt;
diff --git a/gptme/llm/llm_openai.py b/gptme/llm/llm_openai.py
@@ -171,13 +171,6 @@ def _prep_deepseek_reasoner(msgs: list[Message]) -> Generator[Message, None, Non
     yield from _merge_consecutive(_prep_o1(msgs[1:]))
 
 
-def _is_reasoner(base_model: str) -> bool:
-    is_o1 = any(base_model.startswith(om) for om in ["o1", "o3", "o4"])
-    is_deepseek_reasoner = base_model == "deepseek-reasoner"
-    is_gpt5 = base_model.startswith("gpt-5")
-    return is_o1 or is_deepseek_reasoner or is_gpt5
-
-
 @lru_cache(maxsize=2)
 def _is_proxy(client: "OpenAI") -> bool:
     proxy_url = get_config().get_env("LLM_PROXY_URL")
@@ -193,13 +186,15 @@ def chat(messages: list[Message], model: str, tools: list[ToolSpec] | None) -> s
     # top_p controls diversity, temperature controls randomness
 
     from . import _get_base_model, get_provider_from_model  # fmt: skip
+    from .models import get_model  # fmt: skip
 
     provider = get_provider_from_model(model)
     client = get_client(provider)
     is_proxy = _is_proxy(client)
 
     base_model = _get_base_model(model)
-    is_reasoner = _is_reasoner(base_model)
+    model_meta = get_model(model)
+    is_reasoner = model_meta.supports_reasoning
 
     # make the model name prefix with the provider if using LLM_PROXY, to make proxy aware of the provider
     api_model = model if is_proxy else base_model
@@ -216,9 +211,9 @@ def chat(messages: list[Message], model: str, tools: list[ToolSpec] | None) -> s
         top_p=TOP_P if not is_reasoner else NOT_GIVEN,
         tools=tools_dict if tools_dict else NOT_GIVEN,
         extra_headers=extra_headers(provider),
-        extra_body=extra_body(provider, base_model),
+        extra_body=extra_body(provider, model_meta),
     )
-    _record_usage(response.usage, base_model)
+    _record_usage(response.usage, model)
     choice = response.choices[0]
     result = []
     if choice.finish_reason == "tool_calls":
@@ -252,12 +247,14 @@ def extra_headers(provider: Provider) -> dict[str, str]:
     return headers
 
 
-def extra_body(provider: Provider, base_model: str) -> dict[str, Any]:
+def extra_body(provider: Provider, model_meta: ModelMeta) -> dict[str, Any]:
     """Return extra body for the OpenAI API based on the model."""
     body: dict[str, Any] = {}
     if provider == "openrouter":
-        if ":" in base_model:
-            provider_override = base_model.split(":")[1]
+        if model_meta.supports_reasoning:
+            body["reasoning"] = {"enabled": True, "max_tokens": 20000}
+        if "@" in model_meta.model:
+            provider_override = model_meta.model.split("@")[1]
             body["provider"] = {
                 "order": [provider_override],
                 "allow_fallbacks": False,
@@ -269,13 +266,15 @@ def stream(
     messages: list[Message], model: str, tools: list[ToolSpec] | None
 ) -> Generator[str, None, None]:
     from . import _get_base_model, get_provider_from_model  # fmt: skip
+    from .models import get_model  # fmt: skip
 
     provider = get_provider_from_model(model)
     client = get_client(provider)
     is_proxy = _is_proxy(client)
 
     base_model = _get_base_model(model)
-    is_reasoner = _is_reasoner(base_model)
+    model_meta = get_model(model)
+    is_reasoner = model_meta.supports_reasoning
 
     # make the model name prefix with the provider if using LLM_PROXY, to make proxy aware of the provider
     api_model = model if is_proxy else base_model
@@ -294,7 +293,7 @@ def stream(
         stream=True,
         tools=tools_dict if tools_dict else NOT_GIVEN,
         extra_headers=extra_headers(provider),
-        extra_body=extra_body(provider, base_model),
+        extra_body=extra_body(provider, model_meta),
         stream_options={"include_usage": True},
     ):
         from openai.types.chat import ChatCompletionChunk  # fmt: skip
@@ -308,7 +307,7 @@ def stream(
 
         # Record usage if available (typically in final chunk)
         if hasattr(chunk, "usage") and chunk.usage:
-            _record_usage(chunk.usage, base_model)
+            _record_usage(chunk.usage, model)
 
         if not chunk.choices:
             continue
@@ -574,47 +573,51 @@ def get_available_models(provider: Provider) -> list[ModelMeta]:
 def openrouter_model_to_modelmeta(model_data: dict) -> ModelMeta:
     """Convert OpenRouter model data to ModelMeta object."""
     pricing = model_data.get("pricing", {})
+    price_input = float(pricing.get("prompt", 0)) * 1_000_000
+    price_output = float(pricing.get("completion", 0)) * 1_000_000
+    vision = "vision" in model_data.get("architecture", {}).get("modality", "")
+    reasoning = "reasoning" in model_data.get("supported_parameters", [])
+    include_reasoning = "include_reasoning" in model_data.get(
+        "supported_parameters", []
+    )
 
     return ModelMeta(
         provider="openrouter",
         model=model_data.get("id", ""),
         context=model_data.get("context_length", 128_000),
         max_output=model_data.get("max_completion_tokens"),
         supports_streaming=True,  # Most OpenRouter models support streaming
-        supports_vision="vision"
-        in model_data.get("architecture", {}).get("modality", ""),
-        supports_reasoning=False,  # Would need to check model-specific capabilities
-        price_input=float(pricing.get("prompt", 0))
-        * 1_000_000,  # Convert to per-1M tokens
-        price_output=float(pricing.get("completion", 0))
-        * 1_000_000,  # Convert to per-1M tokens
+        supports_vision=vision,
+        supports_reasoning=reasoning and include_reasoning,
+        price_input=price_input,
+        price_output=price_output,
     )
 
 
 def _prepare_messages_for_api(
     messages: list[Message], model: str, tools: list[ToolSpec] | None
 ) -> tuple[Iterable[dict], Iterable["ChatCompletionToolParam"] | None]:
-    from . import _get_base_model  # fmt: skip
     from .models import get_model  # fmt: skip
 
     model_meta = get_model(model)
 
-    is_o1 = _get_base_model(model).startswith("o1")
-    if is_o1:
-        messages = list(_prep_o1(messages))
-
-    # without this, deepseek-chat and reasoner can start outputting gibberish after tool calls
-    # similarly, kimi-k2-instruct doesn't acknowledge tool responses/system messages without it, same with magistral
-    # it probably applies to more models/providers, we should figure out which and perhaps make it default behavior
-    # TODO: it seems to apply to a lot of reasoning models, should maybe be default behavior for all reasoning models?
+    # o1 models need _prep_o1 applied to ALL messages (including first), but no merging
     if any(
-        m in model_meta.model
-        for m in [
-            "deepseek-reasoner",
-            "deepseek-chat",
-            "kimi-k2-instruct",
-            "magistral-medium-2506",
-        ]
+        model_meta.model.startswith(om) for om in ["o1", "o3", "o4"]
+    ) or model_meta.model.startswith("gpt-5"):
+        messages = list(_prep_o1(messages))
+    # other reasoning models use deepseek reasoner prep (first message unchanged, then _prep_o1 on rest)
+    elif (
+        any(
+            m in model_meta.model
+            for m in [
+                "deepseek-reasoner",
+                "deepseek-chat",
+                "kimi-k2",
+                "magistral",
+            ]
+        )
+        or model_meta.supports_reasoning
     ):
         messages = list(_prep_deepseek_reasoner(messages))
 
diff --git a/gptme/llm/models.py b/gptme/llm/models.py
@@ -306,6 +306,20 @@ class _ModelDictMeta(TypedDict):
             "price_output": 0.3,
             "supports_vision": True,
         },
+        "moonshotai/kimi-k2": {
+            "context": 262_144,
+            "max_output": 262_144,
+            "price_input": 0.38,
+            "price_output": 1.52,
+            "supports_vision": True,
+        },
+        "moonshotai/kimi-k2-0905": {
+            "context": 262_144,
+            "max_output": 262_144,
+            "price_input": 0.38,
+            "price_output": 1.52,
+            "supports_vision": True,
+        },
     },
     "nvidia": {},
     "azure": {},
@@ -351,24 +365,58 @@ def get_model(model: str) -> ModelMeta:
         model = get_recommended_model(provider)
         return get_model(f"{provider}/{model}")
 
-    if any(f"{provider}/" in model for provider in PROVIDERS):
-        provider, model = cast(tuple[Provider, str], model.split("/", 1))
-        if provider not in MODELS or model not in MODELS[provider]:
+    # Check if model has provider/model format
+    if any(model.startswith(f"{provider}/") for provider in PROVIDERS):
+        provider_str, model_name = model.split("/", 1)
+
+        # Check if provider is known
+        if provider_str in PROVIDERS:
+            provider = cast(Provider, provider_str)
+
+            # First try static MODELS dict for performance
+            if provider in MODELS and model_name in MODELS[provider]:
+                return ModelMeta(provider, model_name, **MODELS[provider][model_name])
+
+            # For providers that support dynamic fetching, use _get_models_for_provider
+            if provider == "openrouter":
+                try:
+                    models = _get_models_for_provider(provider, dynamic_fetch=True)
+                    for model_meta in models:
+                        if model_meta.model == model_name:
+                            return model_meta
+                except Exception:
+                    # Fall back to unknown model metadata
+                    pass
+
+            # Unknown model, use fallback metadata
             if provider not in ["openrouter", "local"]:
                 log_warn_once(
-                    f"Unknown model: using fallback metadata for {provider}/{model}"
+                    f"Unknown model: using fallback metadata for {provider}/{model_name}"
                 )
-            return ModelMeta(provider, model, context=128_000)
-    else:
-        # try to find model in all providers
-        for provider in MODELS:
-            if model in MODELS[provider]:
-                break
+            return ModelMeta(provider, model_name, context=128_000)
         else:
+            # Unknown provider
             logger.warning(f"Unknown model {model}, using fallback metadata")
             return ModelMeta(provider="unknown", model=model, context=128_000)
+    else:
+        # try to find model in all providers, starting with static models
+        for provider in cast(list[Provider], MODELS.keys()):
+            if model in MODELS[provider]:
+                return ModelMeta(provider, model, **MODELS[provider][model])
+
+        # For model name without provider, also try dynamic fetching for openrouter
+        try:
+            openrouter_models = _get_models_for_provider(
+                "openrouter", dynamic_fetch=True
+            )
+            for model_meta in openrouter_models:
+                if model_meta.model == model:
+                    return model_meta
+        except Exception:
+            pass
 
-    return ModelMeta(provider, model, **MODELS[provider][model])
+        logger.warning(f"Unknown model {model}, using fallback metadata")
+        return ModelMeta(provider="unknown", model=model, context=128_000)
 
 
 def get_recommended_model(provider: Provider) -> str:  # pragma: no cover
diff --git a/tests/test_llm_models.py b/tests/test_llm_models.py
@@ -0,0 +1,116 @@
+from unittest.mock import patch
+
+from gptme.llm.models import (
+    get_model,
+    _get_models_for_provider,
+    ModelMeta,
+)
+
+
+def test_get_static_model():
+    """Test getting a model that exists in static MODELS dict."""
+    model = get_model("openai/gpt-4o")
+    assert model.provider == "openai"
+    assert model.model == "gpt-4o"
+    assert model.context > 0
+
+
+def test_get_model_provider_only():
+    """Test getting recommended model when only provider is given."""
+    model = get_model("openai")
+    assert model.provider == "openai"
+    assert model.model == "gpt-5"  # current recommended model
+
+
+def test_get_model_unknown_provider_model():
+    """Test fallback for unknown provider/model combination."""
+    model = get_model("unknown-provider/unknown-model")
+    assert model.provider == "unknown"
+    assert model.model == "unknown-provider/unknown-model"
+    assert model.context == 128_000  # fallback context
+
+
+def test_get_model_by_name_only():
+    """Test getting model by name only (searches all providers)."""
+    model = get_model("gpt-4o")
+    assert model.provider == "openai"
+    assert model.model == "gpt-4o"
+
+
+def test_get_model_unknown_name_only():
+    """Test fallback for unknown model name without provider."""
+    model = get_model("completely-unknown-model")
+    assert model.provider == "unknown"
+    assert model.model == "completely-unknown-model"
+    assert model.context == 128_000
+
+
+@patch("gptme.llm.models._get_models_for_provider")
+def test_get_model_dynamic_fetch_success(mock_get_models):
+    """Test successful dynamic model fetching for OpenRouter."""
+    # Mock a dynamic model
+    dynamic_model = ModelMeta(
+        provider="openrouter",
+        model="test-dynamic-model",
+        context=100_000,
+        price_input=1.0,
+        price_output=2.0,
+    )
+    mock_get_models.return_value = [dynamic_model]
+
+    model = get_model("openrouter/test-dynamic-model")
+    assert model.provider == "openrouter"
+    assert model.model == "test-dynamic-model"
+    assert model.context == 100_000
+    assert model.price_input == 1.0
+
+    mock_get_models.assert_called_once_with("openrouter", dynamic_fetch=True)
+
+
+@patch("gptme.llm.models._get_models_for_provider")
+def test_get_model_dynamic_fetch_failure(mock_get_models):
+    """Test fallback when dynamic model fetching fails."""
+    mock_get_models.side_effect = Exception("API error")
+
+    model = get_model("openrouter/test-dynamic-model")
+    assert model.provider == "openrouter"
+    assert model.model == "test-dynamic-model"
+    assert model.context == 128_000  # fallback
+
+
+@patch("gptme.llm.models._get_models_for_provider")
+def test_get_model_dynamic_fetch_model_not_found(mock_get_models):
+    """Test fallback when dynamic model is not found in results."""
+    other_model = ModelMeta(provider="openrouter", model="other-model", context=100_000)
+    mock_get_models.return_value = [other_model]
+
+    model = get_model("openrouter/test-dynamic-model")
+    assert model.provider == "openrouter"
+    assert model.model == "test-dynamic-model"
+    assert model.context == 128_000  # fallback
+
+
+def test_get_models_for_provider():
+    """Test getting models for a specific provider."""
+    # Test with static models only
+    openai_models = _get_models_for_provider("openai", dynamic_fetch=False)
+    assert len(openai_models) > 0
+    assert all(m.provider == "openai" for m in openai_models)
+
+
+@patch("gptme.llm.models._get_models_for_provider")
+def test_get_model_name_only_with_dynamic_fetch(mock_get_models):
+    """Test model lookup by name only with dynamic fetching from OpenRouter."""
+    # Mock OpenRouter dynamic model
+    dynamic_model = ModelMeta(
+        provider="openrouter", model="test-model", context=100_000
+    )
+    mock_get_models.return_value = [dynamic_model]
+
+    model = get_model("test-model")
+    assert model.provider == "openrouter"
+    assert model.model == "test-model"
+    assert model.context == 100_000
+
+    # Should have tried OpenRouter dynamic fetch
+    mock_get_models.assert_called_with("openrouter", dynamic_fetch=True)
