From f124bc26632aa2eaf4a2f434d44f9d7fff1f1be0 Mon Sep 17 00:00:00 2001 From: Shahmir Khan Date: Sun, 1 Nov 2026 23:03:24 +0500 Subject: [PATCH 1/6] feat(lm_studio): synchronize loaded model context limits --- litellm/llms/lm_studio/model_info.py | 142 ++++++++++ litellm/router_python_core.py | 12 +- tests/unit/llms/lm_studio/test_model_info.py | 249 ++++++++++++++++++ .../proxy_server/test_routes_model_info.py | 23 +- .../unit/test_router_model_cost_isolation.py | 96 ++++++- 5 files changed, 513 insertions(+), 9 deletions(-) create mode 100644 litellm/llms/lm_studio/model_info.py create mode 100644 tests/unit/llms/lm_studio/test_model_info.py diff --git a/litellm/llms/lm_studio/model_info.py b/litellm/llms/lm_studio/model_info.py new file mode 100644 index 00000000000..128cee66379 --- /dev/null +++ b/litellm/llms/lm_studio/model_info.py @@ -0,0 +1,142 @@ +import hashlib +import json +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType +from typing import Annotated, Final, Protocol, TypeAlias + +import httpx +from pydantic import BeforeValidator, ConfigDict + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS +from litellm.types.llms.base import LiteLLMBaseModel + + +class _MetadataCache(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: object, *, ttl: int) -> None: ... + + +def _positive_limit(value: object) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None + + +_TokenLimit: TypeAlias = Annotated[int | None, BeforeValidator(_positive_limit)] + + +class _InstanceConfig(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + context_length: _TokenLimit = None + + +class _Instance(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + config: _InstanceConfig + + +class _ModelV1(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + key: str + loaded_instances: tuple[_Instance, ...] = () + + +class _ModelsV1(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + models: tuple[_ModelV1, ...] + + def context_length(self, model: str) -> int | None: + instances: Final = tuple(chain.from_iterable(card.loaded_instances for card in self.models)) + exact: Final = tuple(instance.config.context_length for instance in instances if instance.id == model) + matching: Final = next((card.loaded_instances for card in self.models if card.key == model), ()) + contexts: Final = exact or tuple(instance.config.context_length for instance in matching) + if not contexts or any(context is None for context in contexts): + return None + return min(context for context in contexts if context is not None) + + +class _ModelV0(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + state: str + loaded_context_length: _TokenLimit = None + + +class _ModelsV0(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + data: tuple[_ModelV0, ...] + + def context_length(self, model: str) -> int | None: + card: Final = next((card for card in self.data if card.id == model), None) + return card.loaded_context_length if card is not None and card.state == "loaded" else None + + +async def get_lm_studio_model_info( + *, + model: str, + api_base: str, + headers: Mapping[str, str], + client: AsyncHTTPHandler, + cache: _MetadataCache, +) -> Mapping[str, int]: + try: + base: Final = httpx.URL(api_base) + except (ValueError, httpx.InvalidURL): + return MappingProxyType({}) + if base.scheme not in ("http", "https") or not base.host or base.username or base.password: + return MappingProxyType({}) + path: Final = base.path.rstrip("/") + suffix: Final = next((suffix for suffix in ("/api/v1", "/api/v0", "/v1") if path.endswith(suffix)), "") + prefix: Final = path[: -len(suffix)] if suffix else path + url: Final = str(base.copy_with(path=prefix + "/api/v1/models", fragment=None)) + cache_key: Final = ( + "lm_studio_model_info:" + hashlib.sha256(json.dumps((url, sorted(headers.items()))).encode()).hexdigest() + ) + cached: Final[object] = cache.get_cache(cache_key) + if isinstance(cached, (_ModelsV1, _ModelsV0)): + return _token_limits(cached.context_length(model)) + try: + response: Final = await client.get( # pyright: ignore[reportUnknownMemberType] # untyped options + url=url, + headers=dict(headers), + timeout=httpx.Timeout(5.0), + follow_redirects=False, + max_response_bytes=2 * 1024 * 1024, + ) + if response.status_code in (404, 405): + legacy_response: Final = await client.get( # pyright: ignore[reportUnknownMemberType] # untyped options + url=str(base.copy_with(path=prefix + "/api/v0/models", fragment=None)), + headers=dict(headers), + timeout=httpx.Timeout(5.0), + follow_redirects=False, + max_response_bytes=2 * 1024 * 1024, + ) + legacy_response.raise_for_status() + models: Final = _ModelsV0.model_validate_json(legacy_response.content) + cache.set_cache(cache_key, models, ttl=MODEL_INFO_REFRESH_SECONDS) + return _token_limits(models.context_length(model)) + response.raise_for_status() + current_models: Final = _ModelsV1.model_validate_json(response.content) + cache.set_cache(cache_key, current_models, ttl=MODEL_INFO_REFRESH_SECONDS) + return _token_limits(current_models.context_length(model)) + except Exception: # noqa: BLE001 # optional metadata must not interrupt proxy refresh + verbose_logger.debug("Could not discover LM Studio model token limits") + cache.set_cache(cache_key, _ModelsV1(models=()), ttl=60) + return MappingProxyType({}) + + +def _token_limits(context: int | None) -> Mapping[str, int]: + return ( + MappingProxyType({"max_tokens": context, "max_input_tokens": context, "max_output_tokens": context}) + if context is not None + else MappingProxyType({}) + ) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..09fbdb1e365 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -230,6 +230,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( ) from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client +from litellm.llms.lm_studio.model_info import get_lm_studio_model_info from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import ( MODEL_INFO_DISCOVERY_PROVIDERS, @@ -12466,7 +12467,7 @@ class PythonCore(DeprecatedRouterMembers): model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")( model=params.model, litellm_params=params ) - if provider not in MODEL_INFO_DISCOVERY_PROVIDERS: + if provider not in MODEL_INFO_DISCOVERY_PROVIDERS and provider != "lm_studio": return if api_base is None or "*" in model or params.get("use_clientside_credentials"): return @@ -12482,7 +12483,8 @@ class PythonCore(DeprecatedRouterMembers): auth_headers: Final = ( MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({}) ) - limits: Final = await get_openai_compatible_model_info( + discover: Final = get_lm_studio_model_info if provider == "lm_studio" else get_openai_compatible_model_info + limits: Final = await discover( model=model, api_base=api_base, headers=MappingProxyType( @@ -12495,7 +12497,11 @@ class PythonCore(DeprecatedRouterMembers): cache=self._core_cache.in_memory_cache, ) model_id: Final = deployment.model_info.id - if not limits or model_id is None or self._core_get_model_info(model_id) is not raw_deployment: + if ( + (not limits and provider != "lm_studio") + or model_id is None + or self._core_get_model_info(model_id) is not raw_deployment + ): return self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1) self._discovered_model_info_cache.delete_cache(model_id) diff --git a/tests/unit/llms/lm_studio/test_model_info.py b/tests/unit/llms/lm_studio/test_model_info.py new file mode 100644 index 00000000000..e8cb92dd2d0 --- /dev/null +++ b/tests/unit/llms/lm_studio/test_model_info.py @@ -0,0 +1,249 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.lm_studio.model_info import get_lm_studio_model_info +from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS + + +@pytest.mark.parametrize( + ("model", "expected"), + ( + ("org/model", 2048), + ("instance-a", 8192), + ("instance-b", 2048), + ("other", 4096), + ("model", None), + ("missing", None), + ("unloaded", None), + ("incomplete", None), + ), +) +async def test_matches_loaded_instances_without_using_theoretical_limits(model: str, expected: int | None) -> None: + responder: Final = Mock( + return_value=httpx.Response( + 200, + json={ + "models": [ + { + "key": "org/model", + "max_context_length": 131072, + "loaded_instances": [ + {"id": "instance-a", "config": {"context_length": 8192}}, + {"id": "instance-b", "config": {"context_length": 2048}}, + ], + }, + {"key": "other", "loaded_instances": [{"id": "other", "config": {"context_length": 4096}}]}, + {"key": "unloaded", "max_context_length": 131072, "loaded_instances": []}, + { + "key": "incomplete", + "loaded_instances": [ + {"id": "valid", "config": {"context_length": 4096}}, + {"id": "invalid", "config": {}}, + ], + }, + ] + }, + ) + ) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(responder)) as client: + handler.client = client + result: Final = await get_lm_studio_model_info( + model=model, api_base="https://backend.test/v1", headers={}, client=handler, cache=InMemoryCache() + ) + assert result == ( + {} + if expected is None + else {"max_tokens": expected, "max_input_tokens": expected, "max_output_tokens": expected} + ) + + +@pytest.mark.parametrize("base_path", ("", "/v1/", "/api/v1", "/api/v0")) +@pytest.mark.parametrize("status", (404, 405)) +async def test_legacy_fallback_preserves_origin_and_proxy_prefix(base_path: str, status: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + assert request.url.host == "backend.test" + assert request.headers["authorization"] == "Bearer local-key" + if request.url.path == "/tenant/api/v1/models": + return httpx.Response(status) + assert request.url.path == "/tenant/api/v0/models" + return httpx.Response( + 200, + json={ + "data": [ + {"id": "org/model", "state": "loaded", "max_context_length": 32768, "loaded_context_length": 4096}, + {"id": "unloaded", "state": "not-loaded", "max_context_length": 32768}, + ] + }, + ) + + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler.client = client + cache: Final = InMemoryCache() + result: Final = await get_lm_studio_model_info( + model="org/model", + api_base=f"https://backend.test/tenant{base_path}", + headers={"authorization": "Bearer local-key"}, + client=handler, + cache=cache, + ) + assert result == {"max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096} + assert ( + await get_lm_studio_model_info( + model="unloaded", + api_base=f"https://backend.test/tenant{base_path}", + headers={"authorization": "Bearer local-key"}, + client=handler, + cache=cache, + ) + == {} + ) + + +@pytest.mark.parametrize("context", (True, 0, -1, "8192", 1.5, None)) +async def test_invalid_context_is_not_advertised(context: object) -> None: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, + json={ + "models": [ + { + "key": "model", + "max_context_length": 32768, + "loaded_instances": [{"id": "model", "config": {"context_length": context}}], + }, + ] + }, + ) + ) + ) as client: + handler.client = client + assert ( + await get_lm_studio_model_info( + model="model", api_base="https://backend.test", headers={}, client=handler, cache=InMemoryCache() + ) + == {} + ) + + +@pytest.mark.parametrize( + ("response", "requests_per_attempt"), + ( + (httpx.Response(401), 1), + (httpx.Response(503), 1), + (httpx.Response(404), 2), + (httpx.Response(405), 2), + (httpx.Response(302, headers={"location": "https://elsewhere.test"}), 1), + (httpx.Response(200, content=b"not json"), 1), + (httpx.Response(200, json={"models": None}), 1), + (httpx.Response(200, json={"data": []}), 1), + (httpx.ReadTimeout("offline"), 1), + (httpx.ConnectError("offline"), 1), + ), +) +async def test_failures_are_negative_cached_without_redirects_or_unbounded_fallbacks( + response: httpx.Response | Exception, + requests_per_attempt: int, +) -> None: + responder: Final = Mock(side_effect=response if isinstance(response, Exception) else None, return_value=response) + clock: Final = Mock(return_value=0) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(responder), follow_redirects=True) as client: + handler.client = client + cache: Final = InMemoryCache(clock=clock) + for _ in range(2): + assert ( + await get_lm_studio_model_info( + model="model", + api_base="https://backend.test", + headers={"authorization": "secret"}, + client=handler, + cache=cache, + ) + == {} + ) + assert responder.call_count == requests_per_attempt + clock.return_value = 61 + assert ( + await get_lm_studio_model_info( + model="model", + api_base="https://backend.test", + headers={"authorization": "secret"}, + client=handler, + cache=cache, + ) + == {} + ) + assert responder.call_count == 2 * requests_per_attempt + + +async def test_cache_is_scoped_to_server_and_credentials_and_refreshes_context() -> None: + clock: Final = Mock(return_value=0) + contexts: Final = iter((8192, 4096, 2048, 1024)) + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "models": [ + { + "key": "model", + "loaded_instances": [{"id": "instance", "config": {"context_length": next(contexts)}}], + } + ] + }, + ) + + responder: Final = Mock(side_effect=respond) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(responder)) as client: + handler.client = client + cache: Final = InMemoryCache(clock=clock) + + async def lookup(model: str = "model", host: str = "one.test", key: str = "one") -> Mapping[str, int]: + return await get_lm_studio_model_info( + model=model, api_base=f"https://{host}/v1", headers={"authorization": key}, client=handler, cache=cache + ) + + assert (await lookup())["max_input_tokens"] == 8192 + assert (await lookup("instance"))["max_input_tokens"] == 8192 + assert responder.call_count == 1 + assert (await lookup(key="two"))["max_input_tokens"] == 4096 + assert (await lookup(host="two.test"))["max_input_tokens"] == 2048 + clock.return_value = MODEL_INFO_REFRESH_SECONDS + 1 + assert (await lookup())["max_input_tokens"] == 1024 + assert responder.call_count == 4 + + +@pytest.mark.parametrize("api_base", ("", "file:///tmp/server", "https://user:password@backend.test", "http://[:::]")) +async def test_invalid_or_missing_server_never_sends_credentials(api_base: str) -> None: + responder: Final = Mock(side_effect=AssertionError("no network request expected")) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(responder)) as client: + handler.client = client + assert ( + await get_lm_studio_model_info( + model="model", + api_base=api_base, + headers={"authorization": "secret"}, + client=handler, + cache=InMemoryCache(), + ) + == {} + ) + assert responder.call_count == 0 diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index bffa4612964..9bc62a1bc06 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -119,7 +119,9 @@ async def test_discovery_preserves_model_info_fallbacks( _invalidate_model_cost_lowercase_map() +@pytest.mark.parametrize("provider", ("hosted_vllm", "lm_studio")) async def test_upstream_limits_reach_model_info_routes( + provider: str, client: TestClient, auth_as: Callable[[], AbstractContextManager[object]], monkeypatch: pytest.MonkeyPatch, @@ -131,7 +133,7 @@ async def test_upstream_limits_reach_model_info_routes( { "model_name": "local", "litellm_params": { - "model": "hosted_vllm/org/local-model", + "model": f"{provider}/org/local-model", "api_base": "https://backend.test/v1", "api_key": "local-key", }, @@ -144,8 +146,23 @@ async def test_upstream_limits_reach_model_info_routes( monkeypatch.setattr(proxy_server, "user_model", None) def respond(request: httpx.Request) -> httpx.Response: - assert request.url.path == "/v1/models" - return httpx.Response(200, json={"data": [{"id": "org/local-model", "max_model_len": 4096}]}) + assert request.url.path == ("/api/v1/models" if provider == "lm_studio" else "/v1/models") + return httpx.Response( + 200, + json=( + { + "models": [ + { + "key": "org/local-model", + "max_context_length": 32768, + "loaded_instances": [{"id": "org/local-model", "config": {"context_length": 4096}}], + } + ] + } + if provider == "lm_studio" + else {"data": [{"id": "org/local-model", "max_model_len": 4096}]} + ), + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 91f0b2339c7..3d425b6b956 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -31,6 +31,7 @@ import litellm from litellm import Router from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE +from litellm.router_api import ChangeSet from litellm.litellm_core_utils.ptu_pricing import ptu_config_error from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -340,7 +341,7 @@ async def test_discovered_limits_expire_after_the_last_successful_refresh(monkey _invalidate_model_cost_lowercase_map() -@pytest.mark.parametrize("provider", ("hosted_vllm", "openai", "openai_like", "text-completion-openai")) +@pytest.mark.parametrize("provider", ("hosted_vllm", "openai", "openai_like", "text-completion-openai", "lm_studio")) async def test_discovered_limits_are_isolated_overridable_and_refreshable( provider: str, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -348,9 +349,24 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( upstream_limit: Final = iter((8192, 4096, 16384, 2048)) def respond(request: httpx.Request) -> httpx.Response: - assert request.url.path == "/v1/models" + assert request.url.path == ("/api/v1/models" if provider == "lm_studio" else "/v1/models") assert request.headers["authorization"] == "Bearer local-key" - return httpx.Response(200, json={"data": [{"id": "org/local-model", "max_model_len": next(upstream_limit)}]}) + limit: Final = next(upstream_limit) + return httpx.Response( + 200, + json=( + { + "models": [ + { + "key": "org/local-model", + "loaded_instances": [{"id": "org/local-model", "config": {"context_length": limit}}], + } + ] + } + if provider == "lm_studio" + else {"data": [{"id": "org/local-model", "max_model_len": limit}]} + ), + ) router: Final = Router( model_list=[ @@ -4502,3 +4518,77 @@ def test_pricing_provider_fills_unset_rates_from_that_providers_catalog_row() -> finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize( + "response", + ( + httpx.Response(200, json={"models": [{"key": "org/model", "loaded_instances": []}]}), + httpx.Response(200, json={"models": []}), + httpx.Response(503), + ), +) +@pytest.mark.parametrize("input_limit", (None, 1000)) +async def test_lm_studio_refresh_clears_active_limits_but_preserves_explicit_overrides( + response: httpx.Response, + input_limit: int | None, + local_model_cost_map: None, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "alias", + "litellm_params": {"model": "lm_studio/org/model", "api_base": "https://backend.test/v1"}, + "model_info": {"id": "lm-deployment", "max_output_tokens": 123, "max_input_tokens": input_limit}, + } + ] + ) + responses: Final = iter( + ( + httpx.Response( + 200, + json={ + "models": [ + { + "key": "org/model", + "loaded_instances": [{"id": "org/model", "config": {"context_length": 4096}}], + } + ] + }, + ), + response, + ) + ) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: + handler.client = client + await router.aapply(ChangeSet(refresh_model_info=True, refresh_client=handler)) + assert router.registry().group("alias", field="token_limits") == (input_limit or 4096, 123) + router.cache.in_memory_cache.flush_cache() + await router.aapply(ChangeSet(refresh_model_info=True, refresh_client=handler)) + assert router.registry().group("alias", field="token_limits") == (input_limit, 123) + + +async def test_lm_studio_refresh_without_server_does_not_send_api_key( + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + monkeypatch.delenv("LM_STUDIO_API_BASE", raising=False) + monkeypatch.setenv("LM_STUDIO_API_KEY", "local-only-key") + router: Final = Router( + model_list=[ + { + "model_name": "alias", + "litellm_params": {"model": "lm_studio/org/model"}, + } + ] + ) + responder: Final = Mock(side_effect=AssertionError("LM Studio credentials must not reach a default endpoint")) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(responder)) as client: + handler.client = client + await router.aapply(ChangeSet(refresh_model_info=True, refresh_client=handler)) + assert router.registry().group("alias", field="token_limits") == (None, None) + assert responder.call_count == 0 From 95ab955f40c23a2234a0c98625ea5f74e21ee86b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 10:26:50 -0800 Subject: [PATCH 2/6] test: exempt the LM Studio model info cache from the Redis key family check --- tests/unit/test_internal_context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 3336cce9aaf..d428ef4fcde 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -58,6 +58,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/llms/gigachat/authenticator.py", "litellm/llms/github_copilot/per_user_auth.py", "litellm/llms/litellm_proxy/skills/handler.py", + "litellm/llms/lm_studio/model_info.py", "litellm/llms/openai/common_utils.py", "litellm/llms/openai_like/model_info.py", "litellm/llms/vertex_ai/vertex_ai_non_gemini.py", From 286a88bc547aa7d4f27702b7239cc3097520d5cd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 11:43:13 -0800 Subject: [PATCH 3/6] fix(lm_studio): keep last limits on failed discovery, fall back to v0 on unknown-route bodies, refresh each interval --- litellm/llms/lm_studio/model_info.py | 97 ++++++++++++------- litellm/router_python_core.py | 48 +++++++-- tests/unit/llms/lm_studio/test_model_info.py | 69 ++++++++----- .../unit/test_router_model_cost_isolation.py | 75 ++++++++++---- 4 files changed, 200 insertions(+), 89 deletions(-) diff --git a/litellm/llms/lm_studio/model_info.py b/litellm/llms/lm_studio/model_info.py index 128cee66379..bd2a3d5a413 100644 --- a/litellm/llms/lm_studio/model_info.py +++ b/litellm/llms/lm_studio/model_info.py @@ -3,16 +3,17 @@ import json from collections.abc import Mapping from itertools import chain from types import MappingProxyType -from typing import Annotated, Final, Protocol, TypeAlias +from typing import Annotated, Final, Protocol, TypeAlias, TypeVar import httpx -from pydantic import BeforeValidator, ConfigDict +from pydantic import BeforeValidator, ConfigDict, ValidationError from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.types.llms.base import LiteLLMBaseModel +RESPONSE_REUSE_SECONDS: Final = 60 + class _MetadataCache(Protocol): def get_cache(self, key: str) -> object: ... @@ -80,6 +81,13 @@ class _ModelsV0(LiteLLMBaseModel): return card.loaded_context_length if card is not None and card.state == "loaded" else None +class _Unavailable(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + +_Models: Final = TypeVar("_Models", _ModelsV1, _ModelsV0) + + async def get_lm_studio_model_info( *, model: str, @@ -87,51 +95,70 @@ async def get_lm_studio_model_info( headers: Mapping[str, str], client: AsyncHTTPHandler, cache: _MetadataCache, -) -> Mapping[str, int]: +) -> Mapping[str, int] | None: try: base: Final = httpx.URL(api_base) except (ValueError, httpx.InvalidURL): - return MappingProxyType({}) + return None if base.scheme not in ("http", "https") or not base.host or base.username or base.password: - return MappingProxyType({}) + return None path: Final = base.path.rstrip("/") suffix: Final = next((suffix for suffix in ("/api/v1", "/api/v0", "/v1") if path.endswith(suffix)), "") prefix: Final = path[: -len(suffix)] if suffix else path - url: Final = str(base.copy_with(path=prefix + "/api/v1/models", fragment=None)) + current_url: Final = base.copy_with(path=prefix + "/api/v1/models", fragment=None) cache_key: Final = ( - "lm_studio_model_info:" + hashlib.sha256(json.dumps((url, sorted(headers.items()))).encode()).hexdigest() + "lm_studio_model_info:" + + hashlib.sha256(json.dumps((str(current_url), sorted(headers.items()))).encode()).hexdigest() ) cached: Final[object] = cache.get_cache(cache_key) + if isinstance(cached, _Unavailable): + return None if isinstance(cached, (_ModelsV1, _ModelsV0)): return _token_limits(cached.context_length(model)) - try: - response: Final = await client.get( # pyright: ignore[reportUnknownMemberType] # untyped options - url=url, - headers=dict(headers), - timeout=httpx.Timeout(5.0), - follow_redirects=False, - max_response_bytes=2 * 1024 * 1024, - ) - if response.status_code in (404, 405): - legacy_response: Final = await client.get( # pyright: ignore[reportUnknownMemberType] # untyped options - url=str(base.copy_with(path=prefix + "/api/v0/models", fragment=None)), - headers=dict(headers), - timeout=httpx.Timeout(5.0), - follow_redirects=False, - max_response_bytes=2 * 1024 * 1024, - ) - legacy_response.raise_for_status() - models: Final = _ModelsV0.model_validate_json(legacy_response.content) - cache.set_cache(cache_key, models, ttl=MODEL_INFO_REFRESH_SECONDS) - return _token_limits(models.context_length(model)) - response.raise_for_status() - current_models: Final = _ModelsV1.model_validate_json(response.content) - cache.set_cache(cache_key, current_models, ttl=MODEL_INFO_REFRESH_SECONDS) - return _token_limits(current_models.context_length(model)) - except Exception: # noqa: BLE001 # optional metadata must not interrupt proxy refresh + models: Final = await _fetch_models( + client=client, + current_url=current_url, + legacy_url=base.copy_with(path=prefix + "/api/v0/models", fragment=None), + headers=headers, + ) + if models is None: verbose_logger.debug("Could not discover LM Studio model token limits") - cache.set_cache(cache_key, _ModelsV1(models=()), ttl=60) - return MappingProxyType({}) + cache.set_cache(cache_key, _Unavailable(), ttl=RESPONSE_REUSE_SECONDS) + return None + cache.set_cache(cache_key, models, ttl=RESPONSE_REUSE_SECONDS) + return _token_limits(models.context_length(model)) + + +async def _fetch_models( + *, client: AsyncHTTPHandler, current_url: httpx.URL, legacy_url: httpx.URL, headers: Mapping[str, str] +) -> _ModelsV1 | _ModelsV0 | None: + try: + current: Final = await _get(client=client, url=current_url, headers=headers) + current_models: Final = _parse(current, _ModelsV1) + if current_models is not None or not (current.is_success or current.status_code in (404, 405)): + return current_models + return _parse(await _get(client=client, url=legacy_url, headers=headers), _ModelsV0) + except Exception: # noqa: BLE001 # optional metadata must not interrupt proxy refresh + return None + + +async def _get(*, client: AsyncHTTPHandler, url: httpx.URL, headers: Mapping[str, str]) -> httpx.Response: + return await client.get( # pyright: ignore[reportUnknownMemberType] # untyped options + url=str(url), + headers=dict(headers), + timeout=httpx.Timeout(5.0), + follow_redirects=False, + max_response_bytes=2 * 1024 * 1024, + ) + + +def _parse(response: httpx.Response, shape: type[_Models]) -> _Models | None: + if not response.is_success: + return None + try: + return shape.model_validate_json(response.content) + except ValidationError: + return None def _token_limits(context: int | None) -> Mapping[str, int]: diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 09fbdb1e365..dac5bb021c1 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -48,7 +48,20 @@ from datetime import datetime, timezone from functools import lru_cache, partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Generic, + Literal, + NoReturn, + Optional, + Protocol, + TypeAlias, + TypeVar, + Union, + cast, +) import anyio import httpx @@ -672,6 +685,29 @@ _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) _RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None) + + +class _ModelInfoDiscoverer(Protocol): + def __call__( + self, *, model: str, api_base: str, headers: Mapping[str, str], client: AsyncHTTPHandler, cache: InMemoryCache + ) -> Awaitable[Mapping[str, int] | None]: ... + + +async def _discover_openai_compatible_model_info( + *, model: str, api_base: str, headers: Mapping[str, str], client: AsyncHTTPHandler, cache: InMemoryCache +) -> Mapping[str, int] | None: + limits: Final = await get_openai_compatible_model_info( + model=model, api_base=api_base, headers=headers, client=client, cache=cache + ) + return limits or None + + +_MODEL_INFO_DISCOVERERS: Final[Mapping[str, _ModelInfoDiscoverer]] = MappingProxyType( + { + **dict.fromkeys(MODEL_INFO_DISCOVERY_PROVIDERS, _discover_openai_compatible_model_info), + "lm_studio": get_lm_studio_model_info, + } +) _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -12467,7 +12503,8 @@ class PythonCore(DeprecatedRouterMembers): model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")( model=params.model, litellm_params=params ) - if provider not in MODEL_INFO_DISCOVERY_PROVIDERS and provider != "lm_studio": + discover: Final = _MODEL_INFO_DISCOVERERS.get(provider) + if discover is None: return if api_base is None or "*" in model or params.get("use_clientside_credentials"): return @@ -12483,7 +12520,6 @@ class PythonCore(DeprecatedRouterMembers): auth_headers: Final = ( MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({}) ) - discover: Final = get_lm_studio_model_info if provider == "lm_studio" else get_openai_compatible_model_info limits: Final = await discover( model=model, api_base=api_base, @@ -12497,11 +12533,7 @@ class PythonCore(DeprecatedRouterMembers): cache=self._core_cache.in_memory_cache, ) model_id: Final = deployment.model_info.id - if ( - (not limits and provider != "lm_studio") - or model_id is None - or self._core_get_model_info(model_id) is not raw_deployment - ): + if limits is None or model_id is None or self._core_get_model_info(model_id) is not raw_deployment: return self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1) self._discovered_model_info_cache.delete_cache(model_id) diff --git a/tests/unit/llms/lm_studio/test_model_info.py b/tests/unit/llms/lm_studio/test_model_info.py index e8cb92dd2d0..9c8b5fe0d5f 100644 --- a/tests/unit/llms/lm_studio/test_model_info.py +++ b/tests/unit/llms/lm_studio/test_model_info.py @@ -66,13 +66,22 @@ async def test_matches_loaded_instances_without_using_theoretical_limits(model: @pytest.mark.parametrize("base_path", ("", "/v1/", "/api/v1", "/api/v0")) -@pytest.mark.parametrize("status", (404, 405)) -async def test_legacy_fallback_preserves_origin_and_proxy_prefix(base_path: str, status: int) -> None: +@pytest.mark.parametrize( + "current_api_response", + ( + httpx.Response(404), + httpx.Response(405), + httpx.Response(200, json={"error": "Unexpected endpoint or method. (GET /api/v1/models)"}), + ), +) +async def test_legacy_fallback_preserves_origin_and_proxy_prefix( + base_path: str, current_api_response: httpx.Response +) -> None: def respond(request: httpx.Request) -> httpx.Response: assert request.url.host == "backend.test" assert request.headers["authorization"] == "Bearer local-key" if request.url.path == "/tenant/api/v1/models": - return httpx.Response(status) + return current_api_response assert request.url.path == "/tenant/api/v0/models" return httpx.Response( 200, @@ -146,14 +155,13 @@ async def test_invalid_context_is_not_advertised(context: object) -> None: (httpx.Response(404), 2), (httpx.Response(405), 2), (httpx.Response(302, headers={"location": "https://elsewhere.test"}), 1), - (httpx.Response(200, content=b"not json"), 1), - (httpx.Response(200, json={"models": None}), 1), - (httpx.Response(200, json={"data": []}), 1), + (httpx.Response(200, content=b"not json"), 2), + (httpx.Response(200, json={"models": None}), 2), (httpx.ReadTimeout("offline"), 1), (httpx.ConnectError("offline"), 1), ), ) -async def test_failures_are_negative_cached_without_redirects_or_unbounded_fallbacks( +async def test_failures_are_reported_as_unknown_and_negative_cached_without_redirects_or_unbounded_fallbacks( response: httpx.Response | Exception, requests_per_attempt: int, ) -> None: @@ -164,30 +172,37 @@ async def test_failures_are_negative_cached_without_redirects_or_unbounded_fallb async with httpx.AsyncClient(transport=httpx.MockTransport(responder), follow_redirects=True) as client: handler.client = client cache: Final = InMemoryCache(clock=clock) - for _ in range(2): - assert ( - await get_lm_studio_model_info( - model="model", - api_base="https://backend.test", - headers={"authorization": "secret"}, - client=handler, - cache=cache, - ) - == {} - ) - assert responder.call_count == requests_per_attempt - clock.return_value = 61 - assert ( - await get_lm_studio_model_info( + + async def lookup() -> Mapping[str, int] | None: + return await get_lm_studio_model_info( model="model", api_base="https://backend.test", headers={"authorization": "secret"}, client=handler, cache=cache, ) + + assert await lookup() is None + assert await lookup() is None + assert responder.call_count == requests_per_attempt + clock.return_value = 61 + assert await lookup() is None + assert responder.call_count == 2 * requests_per_attempt + + +async def test_server_without_the_model_loaded_is_known_to_have_no_limits() -> None: + responder: Final = Mock(return_value=httpx.Response(200, json={"models": []})) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(responder)) as client: + handler.client = client + assert ( + await get_lm_studio_model_info( + model="model", api_base="https://backend.test", headers={}, client=handler, cache=InMemoryCache() + ) == {} ) - assert responder.call_count == 2 * requests_per_attempt + assert responder.call_count == 1 async def test_cache_is_scoped_to_server_and_credentials_and_refreshes_context() -> None: @@ -215,16 +230,18 @@ async def test_cache_is_scoped_to_server_and_credentials_and_refreshes_context() cache: Final = InMemoryCache(clock=clock) async def lookup(model: str = "model", host: str = "one.test", key: str = "one") -> Mapping[str, int]: - return await get_lm_studio_model_info( + limits: Final = await get_lm_studio_model_info( model=model, api_base=f"https://{host}/v1", headers={"authorization": key}, client=handler, cache=cache ) + assert limits is not None + return limits assert (await lookup())["max_input_tokens"] == 8192 assert (await lookup("instance"))["max_input_tokens"] == 8192 assert responder.call_count == 1 assert (await lookup(key="two"))["max_input_tokens"] == 4096 assert (await lookup(host="two.test"))["max_input_tokens"] == 2048 - clock.return_value = MODEL_INFO_REFRESH_SECONDS + 1 + clock.return_value = MODEL_INFO_REFRESH_SECONDS // 2 assert (await lookup())["max_input_tokens"] == 1024 assert responder.call_count == 4 @@ -244,6 +261,6 @@ async def test_invalid_or_missing_server_never_sends_credentials(api_base: str) client=handler, cache=InMemoryCache(), ) - == {} + is None ) assert responder.call_count == 0 diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 3d425b6b956..dc7b907adb4 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -174,20 +174,26 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch _invalidate_model_cost_lowercase_map() -async def test_discovery_sends_the_secret_behind_an_oidc_api_key(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: +async def test_discovery_sends_the_secret_behind_an_oidc_api_key( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token: Final = tmp_path / "token" token.write_text("local-secret") - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://discovery.test/v1", - "api_key": f"oidc/file/{token}", - }, - "model_info": {"id": "oidc-deployment"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://discovery.test/v1", + "api_key": f"oidc/file/{token}", + }, + "model_info": {"id": "oidc-deployment"}, + } + ] + ) authorizations: Final[list[str | None]] = [] def respond(request: httpx.Request) -> httpx.Response: @@ -293,7 +299,28 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa _invalidate_model_cost_lowercase_map() -async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None: +def _upstream_models_response(provider: str, context: int) -> httpx.Response: + return httpx.Response( + 200, + json=( + { + "models": [ + { + "key": "local-model", + "loaded_instances": [{"id": "local-model", "config": {"context_length": context}}], + } + ] + } + if provider == "lm_studio" + else {"data": [{"id": "local-model", "max_model_len": context}]} + ), + ) + + +@pytest.mark.parametrize("provider", ("hosted_vllm", "lm_studio")) +async def test_discovered_limits_expire_after_the_last_successful_refresh( + provider: str, monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) clock: Final = Mock(return_value=0.0) router: Final = Router( @@ -301,7 +328,7 @@ async def test_discovered_limits_expire_after_the_last_successful_refresh(monkey { "model_name": "local", "litellm_params": { - "model": "hosted_vllm/local-model", + "model": f"{provider}/local-model", "api_base": "https://expiry.test/v1", "api_key": "local-key", }, @@ -312,8 +339,8 @@ async def test_discovered_limits_expire_after_the_last_successful_refresh(monkey router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS) responses: Final = iter( ( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), + _upstream_models_response(provider, 4096), + _upstream_models_response(provider, 8192), httpx.Response(503), ) ) @@ -601,7 +628,9 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) -@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"]) +@pytest.mark.parametrize( + "flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"] +) def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: """A deployment's capability override must not land on the shared ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" @@ -1083,8 +1112,12 @@ def test_bedrock_deployment_keeps_custom_prices_when_profile_lookup_fails( tier_rates: Final = inherit_configured and service_tier != "default" input_rate: Final = configured_prices[f"input_cost_per_token_{service_tier}"] if tier_rates else 0.01 output_rate: Final = configured_prices[f"output_cost_per_token_{service_tier}"] if tier_rates else 0.02 - effective_input: Final = 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate - cache_rate: Final = 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input) + effective_input: Final = ( + 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate + ) + cache_rate: Final = ( + 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input) + ) assert completion_cost( model=model, completion_response=ModelResponse( @@ -2859,7 +2892,8 @@ def test_a_github_copilot_chat_pin_opts_a_dual_endpoint_model_out_of_the_respons @pytest.mark.parametrize("registered_mode, pinned_mode", [("chat", "responses"), ("responses", "chat")]) @pytest.mark.parametrize( - "model, custom_llm_provider", [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")] + "model, custom_llm_provider", + [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")], ) def test_a_github_copilot_responses_pin_still_opts_the_model_into_the_responses_api( monkeypatch: pytest.MonkeyPatch, @@ -4037,7 +4071,9 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info( assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price -def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None: +def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") catalog: Final = litellm.get_model_cost_map(url="") monkeypatch.setattr(litellm, "model_cost", dict(catalog)) @@ -4525,7 +4561,6 @@ def test_pricing_provider_fills_unset_rates_from_that_providers_catalog_row() -> ( httpx.Response(200, json={"models": [{"key": "org/model", "loaded_instances": []}]}), httpx.Response(200, json={"models": []}), - httpx.Response(503), ), ) @pytest.mark.parametrize("input_limit", (None, 1000)) From 3f114526f4b263356c2d9171bf41db7a02a21f64 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:35:35 -0800 Subject: [PATCH 4/6] fix(proxy): keep discovered limits out of PATCH /model/{id}/update echoes --- .../model_management_endpoints.py | 16 ++++++- .../test_model_management_endpoints.py | 45 +++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 2c7453ebb1c..dfb5b05229a 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1420,6 +1420,7 @@ def update_db_model( db_model: Deployment, updated_patch: updateDeployment, loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, + discovered_model_info: Mapping[str, object] = MappingProxyType({}), ) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -1453,6 +1454,7 @@ def update_db_model( incoming_model_info, _cost_map_entry(db_model, incoming_model_info), _loaded_catalog_entry(incoming_model_info, loaded_catalog), + discovered_model_info, ) stored_model_info_fields: Final[Mapping[str, object]] = without_server_derived_pricing(incoming_model_info) merged_model_info.update( @@ -1800,6 +1802,11 @@ async def patch_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, write_row=write_row, + discovered_model_info=( + llm_router.registry().discovered_model_info(model_id) + if llm_router is not None + else MappingProxyType({}) + ), ) if updated_model is None: @@ -2155,6 +2162,7 @@ async def _update_team_model_in_db( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], + discovered_model_info: Mapping[str, object] = MappingProxyType({}), ) -> _RowT: """ Handle team model updates with proper alias management. @@ -2199,7 +2207,9 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + return await write_row( + update_db_model(db_model=db_model, updated_patch=patch_data, discovered_model_info=discovered_model_info) + ) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -2220,7 +2230,9 @@ async def _update_team_model_in_db( ), } ) - row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=team_patch)) + row: Final = await write_row( + update_db_model(db_model=db_model, updated_patch=team_patch, discovered_model_info=discovered_model_info) + ) if is_new_team_assignment: await _setup_new_team_model_assignment( diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 44c43713438..23871dd1f41 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1964,6 +1964,51 @@ class TestUpdateModel: assert persisted["max_output_tokens"] == 64000 assert not {"max_input_tokens", "direct_access", "access_via_team_ids"} & persisted.keys() + @pytest.mark.asyncio + async def test_patch_model_does_not_pin_discovered_limits_echoed_from_model_info(self): + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + db: Final = _ContendedModelDb( + LiteLLM_ProxyModelTable( + model_id="m1", + model_name="my-model", + litellm_params=dict(_LOCAL_DEPLOYMENT_PARAMS), + model_info={"id": "m1"}, + ), + (), + ) + router: Final = Router( + model_list=[ + {"model_name": "my-model", "litellm_params": dict(_LOCAL_DEPLOYMENT_PARAMS), "model_info": {"id": "m1"}} + ] + ) + await _discover_from_upstream(router, {"id": "local-model", "max_model_len": 128000}) + discovered: Final = dict(router.registry().discovered_model_info("m1")) + assert discovered + + with ( + patch("litellm.proxy.proxy_server.prisma_client", db), + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( # test-quality-ok: [TQ008] the route reloads the proxy-global router with no seam to inject + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await patch_model( + model_id="m1", + patch_data=updateDeployment( + model_info=ModelInfo(id="m1", key="hosted_vllm/local-model", tier="paid", **discovered) + ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["tier"] == "paid" + assert not discovered.keys() & persisted.keys() + @pytest.mark.asyncio async def test_update_model_does_not_copy_litellm_params_pricing_into_model_info(self): db: Final = await self._update_model( From a96c2a9af126c51ba68310d05f4167b086d013e9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:09:04 -0800 Subject: [PATCH 5/6] fix(proxy): keep discovered limits out of the model editor view and off LM Studio embedding models The dashboard edit form now gets null for upstream-discovered token limits, so saving it never pins them as explicit overrides on another worker or after an LM Studio reload. LM Studio embedding models no longer advertise their loaded context, since LM Studio truncates over-window embedding input itself and pre-call checks would turn those 200s into 400s. --- litellm/llms/lm_studio/model_info.py | 13 ++- litellm/proxy/proxy_server.py | 22 ++++- tests/unit/llms/lm_studio/test_model_info.py | 59 +++++++++++++ .../proxy_server/test_routes_model_info.py | 82 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 5 files changed, 171 insertions(+), 10 deletions(-) diff --git a/litellm/llms/lm_studio/model_info.py b/litellm/llms/lm_studio/model_info.py index bd2a3d5a413..24ab1f33605 100644 --- a/litellm/llms/lm_studio/model_info.py +++ b/litellm/llms/lm_studio/model_info.py @@ -6,7 +6,7 @@ from types import MappingProxyType from typing import Annotated, Final, Protocol, TypeAlias, TypeVar import httpx -from pydantic import BeforeValidator, ConfigDict, ValidationError +from pydantic import BeforeValidator, ConfigDict, Field, ValidationError from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -45,6 +45,7 @@ class _ModelV1(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) key: str + kind: str | None = Field(default=None, alias="type") loaded_instances: tuple[_Instance, ...] = () @@ -54,9 +55,10 @@ class _ModelsV1(LiteLLMBaseModel): models: tuple[_ModelV1, ...] def context_length(self, model: str) -> int | None: - instances: Final = tuple(chain.from_iterable(card.loaded_instances for card in self.models)) + chat_cards: Final = tuple(card for card in self.models if card.kind != "embedding") + instances: Final = tuple(chain.from_iterable(card.loaded_instances for card in chat_cards)) exact: Final = tuple(instance.config.context_length for instance in instances if instance.id == model) - matching: Final = next((card.loaded_instances for card in self.models if card.key == model), ()) + matching: Final = next((card.loaded_instances for card in chat_cards if card.key == model), ()) contexts: Final = exact or tuple(instance.config.context_length for instance in matching) if not contexts or any(context is None for context in contexts): return None @@ -67,6 +69,7 @@ class _ModelV0(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) id: str + kind: str | None = Field(default=None, alias="type") state: str loaded_context_length: _TokenLimit = None @@ -78,7 +81,9 @@ class _ModelsV0(LiteLLMBaseModel): def context_length(self, model: str) -> int | None: card: Final = next((card for card in self.data if card.id == model), None) - return card.loaded_context_length if card is not None and card.state == "loaded" else None + if card is None or card.state != "loaded" or card.kind == "embeddings": + return None + return card.loaded_context_length class _Unavailable(LiteLLMBaseModel): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 87bc78730b5..0f580b9f621 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17492,7 +17492,10 @@ async def _find_model_by_id( async def _models_with_stored_param_keys( - models: Sequence[Mapping[str, object]], model_id: str, prisma_client: PrismaClient + models: Sequence[Mapping[str, object]], + model_id: str, + prisma_client: PrismaClient, + discovered_model_info: Mapping[str, object], ) -> Sequence[Mapping[str, object]]: visible_model: Final = next( (model for model in models if _with_get(model.get("model_info", {})).get("id") == model_id), None @@ -17511,10 +17514,12 @@ async def _models_with_stored_param_keys( else _LITELLM_PARAMS_COPY.validate_python(raw_params) ) visible_params: Final = _LITELLM_PARAMS_COPY.validate_python(visible_model.get("litellm_params", {})) + visible_info: Final = _LITELLM_PARAMS_COPY.validate_python(visible_model.get("model_info", {})) editor_model: Final = MappingProxyType( { **visible_model, "litellm_params": {key: value for key, value in visible_params.items() if key in stored_params}, + "model_info": {key: None if key in discovered_model_info else value for key, value in visible_info.items()}, } ) return tuple(editor_model if model is visible_model else model for model in models) @@ -17539,7 +17544,10 @@ async def model_info_v2( modelId: str | None = fastapi.Query(None, description="Search for a specific model by its unique ID"), use_stored_params: bool = fastapi.Query( False, - description="With modelId, omit runtime-only litellm_params defaults from a DB-stored model's editor view.", + description=( + "With modelId, omit runtime-only litellm_params defaults and upstream-discovered token limits " + "from a DB-stored model's editor view." + ), ), teamId: str | None = fastapi.Query( None, @@ -17592,7 +17600,8 @@ async def model_info_v2( page / size: Pagination controls (defaults: page=1, size=50). search: Case-insensitive partial match on model name or team public name. modelId: Return a single deployment by LiteLLM model id. - use_stored_params: With modelId, return only persisted litellm_params keys, retaining normal redaction. + use_stored_params: With modelId, return only persisted litellm_params keys and null upstream-discovered + token limits, retaining normal redaction. teamId: Filter to models with direct access or team membership for this team id. sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. blocked: Filter by routing status (false = active, true = paused). @@ -17769,7 +17778,12 @@ async def model_info_v2( if use_stored_params is True and isinstance(modelId, str): return _paginate_models_response( - all_models=await _models_with_stored_param_keys(all_models, modelId, prisma_client), + all_models=await _models_with_stored_param_keys( + all_models, + modelId, + prisma_client, + llm_router.registry().discovered_model_info(modelId), + ), page=page, size=size, total_count=search_total_count, diff --git a/tests/unit/llms/lm_studio/test_model_info.py b/tests/unit/llms/lm_studio/test_model_info.py index 9c8b5fe0d5f..4f5f132de49 100644 --- a/tests/unit/llms/lm_studio/test_model_info.py +++ b/tests/unit/llms/lm_studio/test_model_info.py @@ -118,6 +118,65 @@ async def test_legacy_fallback_preserves_origin_and_proxy_prefix( ) +@pytest.mark.parametrize( + ("path", "models"), + ( + ( + "/api/v1/models", + { + "models": [ + { + "type": "embedding", + "key": "org/embedder", + "loaded_instances": [ + {"id": "org/embedder", "config": {"context_length": 2048}}, + {"id": "embedder-copy", "config": {"context_length": 2048}}, + ], + }, + { + "type": "llm", + "key": "org/chat", + "loaded_instances": [{"id": "org/chat", "config": {"context_length": 4096}}], + }, + ] + }, + ), + ( + "/api/v0/models", + { + "data": [ + {"id": "org/embedder", "type": "embeddings", "state": "loaded", "loaded_context_length": 2048}, + {"id": "embedder-copy", "type": "embeddings", "state": "loaded", "loaded_context_length": 2048}, + {"id": "org/chat", "type": "llm", "state": "loaded", "loaded_context_length": 4096}, + ] + }, + ), + ), +) +async def test_embedding_models_advertise_no_limits_since_lm_studio_truncates_their_input( + path: str, models: Mapping[str, object] +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=models) if request.url.path == path else httpx.Response(404) + + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler.client = client + cache: Final = InMemoryCache() + limits: Final = { + model: await get_lm_studio_model_info( + model=model, api_base="https://backend.test/v1", headers={}, client=handler, cache=cache + ) + for model in ("org/embedder", "embedder-copy", "org/chat") + } + assert limits == { + "org/embedder": {}, + "embedder-copy": {}, + "org/chat": {"max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096}, + } + + @pytest.mark.parametrize("context", (True, 0, -1, "8192", 1.5, None)) async def test_invalid_context_is_not_advertised(context: object) -> None: handler: Final = AsyncHTTPHandler() diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 6f032e8d0c8..16e74752d6f 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -183,6 +183,88 @@ async def test_upstream_limits_reach_model_info_routes( _invalidate_model_cost_lowercase_map() +@pytest.mark.parametrize("provider", ("hosted_vllm", "lm_studio")) +async def test_editor_view_save_never_pins_discovered_limits( + provider: str, + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + mock_prisma: MagicMock, +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + stored_params: Final = { + "model": f"{provider}/org/local-model", + "api_base": "https://backend.test/v1", + "api_key": "local-key", + } + stored_info: Final = {"id": "local-deployment", "db_model": True, "max_output_tokens": 512} + router: Final = litellm.Router( + model_list=[{"model_name": "local", "litellm_params": stored_params, "model_info": stored_info}] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + mock_prisma.db.litellm_proxymodeltable.find_unique.return_value = LiteLLM_ProxyModelTable( + model_id="local-deployment", model_name="local", litellm_params=stored_params, model_info=stored_info + ) + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json=( + { + "models": [ + { + "key": "org/local-model", + "max_context_length": 32768, + "loaded_instances": [{"id": "org/local-model", "config": {"context_length": 4096}}], + } + ] + } + if provider == "lm_studio" + else {"data": [{"id": "org/local-model", "max_model_len": 4096}]} + ), + ) + + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as upstream: + handler.client = upstream + litellm.in_memory_llm_clients_cache.set_cache("async_httpx_clientopenai", handler) + await proxy_server.ProxyStartupEvent.refresh_model_info() + with auth_as(): + shown: Final = client.get("/v2/model/info", params={"modelId": "local-deployment"}) + editor: Final = client.get( + "/v2/model/info", params={"modelId": "local-deployment", "use_stored_params": "true"} + ) + _invalidate_model_cost_lowercase_map() + + assert shown.status_code == editor.status_code == 200, editor.text + discovered: Final = router.registry().discovered_model_info("local-deployment") + assert discovered.get("max_input_tokens") == 4096 + assert shown.json()["data"][0]["model_info"]["max_input_tokens"] == 4096 + editor_info: Final = editor.json()["data"][0]["model_info"] + assert editor_info["max_output_tokens"] == 512 + assert all(editor_info.get(key) is None for key in discovered) + saved_on_a_worker_without_discovery: Final = update_db_model( + db_model=Deployment( + model_name="local", + litellm_params=LiteLLM_Params.model_validate(stored_params), + model_info=ModelInfo.model_validate(stored_info), + ), + updated_patch=PatchModelRequest(model_info=ModelInfo.model_validate(editor_info)), + ) + saved_info: Final = json.loads(saved_on_a_worker_without_discovery["model_info"]) + assert saved_info["max_output_tokens"] == 512 + assert discovered.keys().isdisjoint(saved_info) + + # --------------------------------------------------------------------------- # GET /v2/model/info # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2191454a2f7..7b4059bd960 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24127,7 +24127,8 @@ export interface paths { * page / size: Pagination controls (defaults: page=1, size=50). * search: Case-insensitive partial match on model name or team public name. * modelId: Return a single deployment by LiteLLM model id. - * use_stored_params: With modelId, return only persisted litellm_params keys, retaining normal redaction. + * use_stored_params: With modelId, return only persisted litellm_params keys and null upstream-discovered + * token limits, retaining normal redaction. * teamId: Filter to models with direct access or team membership for this team id. * sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. * blocked: Filter by routing status (false = active, true = paused). @@ -85924,7 +85925,7 @@ export interface operations { search?: string | null; /** @description Search for a specific model by its unique ID */ modelId?: string | null; - /** @description With modelId, omit runtime-only litellm_params defaults from a DB-stored model's editor view. */ + /** @description With modelId, omit runtime-only litellm_params defaults and upstream-discovered token limits from a DB-stored model's editor view. */ use_stored_params?: boolean; /** @description Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids */ teamId?: string | null; From 052b10e805b51963b3f529b55971ed712168b863 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:48:36 -0800 Subject: [PATCH 6/6] fix(lm_studio): read the window of the copy LM Studio serves a model name on When several copies of a model are loaded and none is named after the model, LM Studio sends a request by model name to the first loaded copy. Reading the smallest window instead made pre-call checks refuse prompts that copy serves. --- litellm/llms/lm_studio/model_info.py | 12 ++++++------ tests/unit/llms/lm_studio/test_model_info.py | 12 ++++++++++-- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/litellm/llms/lm_studio/model_info.py b/litellm/llms/lm_studio/model_info.py index 24ab1f33605..e2b72638c68 100644 --- a/litellm/llms/lm_studio/model_info.py +++ b/litellm/llms/lm_studio/model_info.py @@ -57,12 +57,12 @@ class _ModelsV1(LiteLLMBaseModel): def context_length(self, model: str) -> int | None: chat_cards: Final = tuple(card for card in self.models if card.kind != "embedding") instances: Final = tuple(chain.from_iterable(card.loaded_instances for card in chat_cards)) - exact: Final = tuple(instance.config.context_length for instance in instances if instance.id == model) - matching: Final = next((card.loaded_instances for card in chat_cards if card.key == model), ()) - contexts: Final = exact or tuple(instance.config.context_length for instance in matching) - if not contexts or any(context is None for context in contexts): - return None - return min(context for context in contexts if context is not None) + named: Final = next((instance for instance in instances if instance.id == model), None) + first_loaded: Final = next( + (card.loaded_instances[0] for card in chat_cards if card.key == model and card.loaded_instances), None + ) + serving: Final = named if named is not None else first_loaded + return serving.config.context_length if serving is not None else None class _ModelV0(LiteLLMBaseModel): diff --git a/tests/unit/llms/lm_studio/test_model_info.py b/tests/unit/llms/lm_studio/test_model_info.py index 4f5f132de49..c8cb04592d3 100644 --- a/tests/unit/llms/lm_studio/test_model_info.py +++ b/tests/unit/llms/lm_studio/test_model_info.py @@ -14,9 +14,10 @@ from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS @pytest.mark.parametrize( ("model", "expected"), ( - ("org/model", 2048), + ("org/model", 8192), ("instance-a", 8192), ("instance-b", 2048), + ("small-first", 1024), ("other", 4096), ("model", None), ("missing", None), @@ -38,13 +39,20 @@ async def test_matches_loaded_instances_without_using_theoretical_limits(model: {"id": "instance-b", "config": {"context_length": 2048}}, ], }, + { + "key": "small-first", + "loaded_instances": [ + {"id": "small-a", "config": {"context_length": 1024}}, + {"id": "big-b", "config": {"context_length": 4096}}, + ], + }, {"key": "other", "loaded_instances": [{"id": "other", "config": {"context_length": 4096}}]}, {"key": "unloaded", "max_context_length": 131072, "loaded_instances": []}, { "key": "incomplete", "loaded_instances": [ - {"id": "valid", "config": {"context_length": 4096}}, {"id": "invalid", "config": {}}, + {"id": "valid", "config": {"context_length": 4096}}, ], }, ]