diff --git a/litellm/llms/lm_studio/model_info.py b/litellm/llms/lm_studio/model_info.py new file mode 100644 index 00000000000..e2b72638c68 --- /dev/null +++ b/litellm/llms/lm_studio/model_info.py @@ -0,0 +1,174 @@ +import hashlib +import json +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType +from typing import Annotated, Final, Protocol, TypeAlias, TypeVar + +import httpx +from pydantic import BeforeValidator, ConfigDict, Field, ValidationError + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.llms.base import LiteLLMBaseModel + +RESPONSE_REUSE_SECONDS: Final = 60 + + +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 + kind: str | None = Field(default=None, alias="type") + loaded_instances: tuple[_Instance, ...] = () + + +class _ModelsV1(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + models: tuple[_ModelV1, ...] + + 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)) + 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): + model_config = ConfigDict(frozen=True) + + id: str + kind: str | None = Field(default=None, alias="type") + 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) + if card is None or card.state != "loaded" or card.kind == "embeddings": + return None + return card.loaded_context_length + + +class _Unavailable(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + +_Models: Final = TypeVar("_Models", _ModelsV1, _ModelsV0) + + +async def get_lm_studio_model_info( + *, + model: str, + api_base: str, + headers: Mapping[str, str], + client: AsyncHTTPHandler, + cache: _MetadataCache, +) -> Mapping[str, int] | None: + try: + base: Final = httpx.URL(api_base) + except (ValueError, httpx.InvalidURL): + return None + if base.scheme not in ("http", "https") or not base.host or base.username or base.password: + 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 + current_url: Final = base.copy_with(path=prefix + "/api/v1/models", fragment=None) + cache_key: Final = ( + "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)) + 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, _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]: + return ( + MappingProxyType({"max_tokens": context, "max_input_tokens": context, "max_output_tokens": context}) + if context is not None + else MappingProxyType({}) + ) 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/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/litellm/router_python_core.py b/litellm/router_python_core.py index c0c5d74a88a..5305d75b9e7 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 @@ -234,6 +247,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, @@ -684,6 +698,29 @@ _CALL_KWARGS_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]) @@ -12654,7 +12691,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: + 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 @@ -12670,7 +12708,7 @@ 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( + limits: Final = await discover( model=model, api_base=api_base, headers=MappingProxyType( @@ -12683,7 +12721,7 @@ 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 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 new file mode 100644 index 00000000000..c8cb04592d3 --- /dev/null +++ b/tests/unit/llms/lm_studio/test_model_info.py @@ -0,0 +1,333 @@ +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", 8192), + ("instance-a", 8192), + ("instance-b", 2048), + ("small-first", 1024), + ("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": "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": "invalid", "config": {}}, + {"id": "valid", "config": {"context_length": 4096}}, + ], + }, + ] + }, + ) + ) + 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( + "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 current_api_response + 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( + ("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() + 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"), 2), + (httpx.Response(200, json={"models": None}), 2), + (httpx.ReadTimeout("offline"), 1), + (httpx.ConnectError("offline"), 1), + ), +) +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: + 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) + + 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 == 1 + + +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]: + 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 // 2 + 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(), + ) + is None + ) + assert responder.call_count == 0 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( 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 6ad5ef1a4eb..16e74752d6f 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -120,7 +120,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, @@ -132,7 +134,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", }, @@ -145,8 +147,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() @@ -166,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/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index ceb9ef38887..ffe39d73fb0 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", diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 931fabd85ad..35d1e6c3d19 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 @@ -293,7 +294,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 +323,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 +334,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), ) ) @@ -341,7 +363,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: @@ -349,9 +371,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=[ @@ -4662,3 +4699,76 @@ 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": []}), + ), +) +@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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 15d440c0ad4..7ae1b380514 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;