From bb9a71c2850bef0701e28031da3479932e086660 Mon Sep 17 00:00:00 2001 From: NewtonChutney <70827815+NewtonChutney@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:51:26 +0530 Subject: [PATCH] feat(lm_studio): support check_provider_endpoint model discovery get_provider_model_info() had no LM_STUDIO branch, so /v1/models with check_provider_endpoint=true silently returned the literal wildcard string (e.g. "lm_studio/*") instead of the models actually loaded on the LM Studio server. lm_studio was also missing from models_by_provider, so the proxy's wildcard-model expansion (get_provider_models) bailed out before ever reaching get_valid_models. - utils.py: add LM_STUDIO to get_provider_model_info's dispatch table; refactor the no-import/no-model-arg branches into a dict lookup instead of adding a 16th elif, to stay under the C901 complexity budget this file is already at - __init__.py: register lm_studio in models_by_provider (empty set, same pattern as lemonade -- no fixed catalog, discovery-only) - lm_studio/chat/transformation.py: LMStudioChatConfig.get_models() resolves api_base/api_key via the same LM_STUDIO_API_BASE/KEY fallback chat completions already use (inherited OpenAIGPTConfig default points at api.openai.com, which is wrong here), and prefixes discovered ids with "lm_studio/" matching OllamaModelInfo's convention Verified against a live LM Studio instance over Tailscale: a stock (unpatched) no-DB proxy's /v1/models returns only the literal "lm_studio/*"; with this fix it returns the 9 real models actually loaded (qwen/qwen3.6-35b-a3b, meta/muse-glimmer, etc.), each correctly prefixed and independently callable. --- litellm/__init__.py | 2 + litellm/llms/lm_studio/chat/transformation.py | 19 +++++++ litellm/utils.py | 39 +++++++------- tests/test_litellm/test_utils.py | 52 +++++++++++++++++++ 4 files changed, 94 insertions(+), 18 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 056dd532f5f..154de8b330c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -661,6 +661,7 @@ wandb_models: Set = set(WANDB_MODELS) ovhcloud_models: Set = set() ovhcloud_embedding_models: Set = set() lemonade_models: Set = set() +lm_studio_models: Set = set() # mutable-ok: dynamic-discovery-only, like lemonade_models docker_model_runner_models: Set = set() amazon_nova_models: Set = set() stability_models: Set = set() @@ -1177,6 +1178,7 @@ def _build_models_by_provider() -> dict: "wandb": wandb_models, "ovhcloud": ovhcloud_models | ovhcloud_embedding_models, "lemonade": lemonade_models, + "lm_studio": lm_studio_models, "clarifai": clarifai_models, "amazon_nova": amazon_nova_models, "stability": stability_models, diff --git a/litellm/llms/lm_studio/chat/transformation.py b/litellm/llms/lm_studio/chat/transformation.py index 54a73bdc053..f440e1685c4 100644 --- a/litellm/llms/lm_studio/chat/transformation.py +++ b/litellm/llms/lm_studio/chat/transformation.py @@ -19,6 +19,25 @@ class LMStudioChatConfig(OpenAIGPTConfig): ) # LM Studio does not require an api key, but OpenAI client requires non-None value return api_base, dynamic_api_key + def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: + """ + Calls LM Studio's `/v1/models` endpoint and returns the list of models, + prefixed with "lm_studio/" (matching the OllamaModelInfo convention) + so discovered ids are directly callable without the caller having to + add the provider prefix themselves. + + Reuses the same api_base/api_key resolution as chat completions + (LM_STUDIO_API_BASE / LM_STUDIO_API_KEY env vars, "fake-api-key" + fallback) instead of OpenAIGPTConfig's default of + https://api.openai.com, which would be wrong for a local/self-hosted + LM Studio server. + """ + api_base, api_key = self._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) + models: Final = super().get_models(api_key=api_key, api_base=api_base) + return [ # mutable-ok: matches OllamaModelInfo.get_models' list-of-prefixed-ids contract + m if m.startswith("lm_studio/") else f"lm_studio/{m}" for m in models + ] + def map_openai_params( self, non_default_params: dict, diff --git a/litellm/utils.py b/litellm/utils.py index 79372f00284..03f5a19518e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8522,24 +8522,31 @@ class ProviderConfigManager: model: str | None, provider: LlmProviders, ) -> BaseLLMModelInfo | None: - if LlmProviders.FIREWORKS_AI == provider: - return litellm.FireworksAIConfig() - elif LlmProviders.OPENAI == provider: - return litellm.OpenAIGPTConfig() - elif LlmProviders.GEMINI == provider: - return litellm.GeminiModelInfo() + # Providers whose BaseLLMModelInfo needs no local import and no `model` + # arg: a single dict lookup instead of N `elif` branches keeps this + # dispatcher's cyclomatic complexity from growing with every provider + # added here (see the elif chain below for providers that DO need a + # local import or the `model` arg). + simple_provider_configs: Final[ + dict[LlmProviders, Callable[[], BaseLLMModelInfo]] + ] = { # mutable-ok: built once per call, read-only lookup table replacing 9 trivial elif branches + LlmProviders.FIREWORKS_AI: litellm.FireworksAIConfig, + LlmProviders.OPENAI: litellm.OpenAIGPTConfig, + LlmProviders.GEMINI: litellm.GeminiModelInfo, + LlmProviders.LITELLM_PROXY: litellm.LiteLLMProxyChatConfig, + LlmProviders.TOPAZ: litellm.TopazModelInfo, + LlmProviders.ANTHROPIC: litellm.AnthropicModelInfo, + LlmProviders.XAI: litellm.XAIModelInfo, + LlmProviders.LEMONADE: litellm.LemonadeChatConfig, + LlmProviders.CLARIFAI: litellm.ClarifaiConfig, + LlmProviders.LM_STUDIO: litellm.LMStudioChatConfig, + } + if provider in simple_provider_configs: + return simple_provider_configs[provider]() elif LlmProviders.VERTEX_AI == provider: from litellm.llms.vertex_ai.common_utils import VertexAIModelInfo return VertexAIModelInfo() - elif LlmProviders.LITELLM_PROXY == provider: - return litellm.LiteLLMProxyChatConfig() - elif LlmProviders.TOPAZ == provider: - return litellm.TopazModelInfo() - elif LlmProviders.ANTHROPIC == provider: - return litellm.AnthropicModelInfo() - elif LlmProviders.XAI == provider: - return litellm.XAIModelInfo() elif LlmProviders.OLLAMA == provider or LlmProviders.OLLAMA_CHAT == provider: # Dynamic model listing for Ollama server from litellm.llms.ollama.common_utils import OllamaModelInfo @@ -8551,10 +8558,6 @@ class ProviderConfigManager: ) return VLLMModelInfo() - elif LlmProviders.LEMONADE == provider: - return litellm.LemonadeChatConfig() - elif LlmProviders.CLARIFAI == provider: - return litellm.ClarifaiConfig() elif LlmProviders.BEDROCK == provider: from litellm.llms.bedrock.common_utils import BedrockModelInfo diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e9e6167fb9..414074e9d26 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2814,6 +2814,58 @@ class TestGetValidModelsWithCLI: assert headers.get("Authorization") == "Bearer sk-test-cli-key-123" +class TestGetValidModelsLMStudio: + """Test get_valid_models(check_provider_endpoint=True) for lm_studio. + + get_provider_model_info() previously had no LM_STUDIO branch, so + discovery silently returned an empty list regardless of + check_provider_endpoint; and lm_studio was missing from + models_by_provider, so the proxy's wildcard-model expansion + (get_provider_models) bailed out before ever calling get_valid_models. + """ + + def test_get_valid_models_lm_studio_discovery(self): + """Discovery hits LM Studio's own /v1/models, not OpenAI's default api_base, and prefixes results with lm_studio/.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "data": [ + {"id": "qwen/qwen3-a3b", "object": "model"}, + {"id": "lm_studio/already-prefixed", "object": "model"}, + ] + } + + with patch.object( + litellm.module_level_client, "get", return_value=mock_response + ) as mock_get: + result = litellm.get_valid_models( + check_provider_endpoint=True, + custom_llm_provider="lm_studio", + api_key="lm-studio-key", + api_base="http://my-lm-studio-host:1234", + ) + + assert isinstance(result, list) + assert result == [ + "lm_studio/qwen/qwen3-a3b", + "lm_studio/already-prefixed", + ] + + mock_get.assert_called_once() + _, call_kwargs = mock_get.call_args + + # Must hit the passed-in LM Studio host, not OpenAIGPTConfig's + # default of https://api.openai.com + assert call_kwargs["url"] == "http://my-lm-studio-host:1234/v1/models" + assert call_kwargs["headers"]["Authorization"] == "Bearer lm-studio-key" + + def test_lm_studio_in_models_by_provider(self): + """lm_studio must be a key in models_by_provider or the proxy's + wildcard-model expansion (get_provider_models) bails out before + ever reaching get_valid_models/check_provider_endpoint.""" + assert "lm_studio" in litellm.models_by_provider + + class TestIsCachedMessage: """Test is_cached_message function for context caching detection.