feat(lm_studio): synchronize loaded model context limits

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({})
+    )

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)

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

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()

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