From f124bc26632aa2eaf4a2f434d44f9d7fff1f1be0 Mon Sep 17 00:00:00 2001 From: Shahmir Khan Date: Sun, 1 Nov 2026 23:03:24 +0500 Subject: [PATCH] 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