scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

test(proxy): inject the loaded catalog into update_db_model instead of patching the class
Some checks are pending
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Details

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-20 10:11:02 +00:00
parent 65849225f6
commit 0141ea4a31
2 changed files with 16 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -19,7 +19,7 @@ from dataclasses import dataclass
from fnmatch import fnmatchcase
from json import JSONDecodeError
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
@ -888,14 +888,23 @@ def _cost_map_entry(db_model: Deployment, incoming_model_info: Mapping[str, obje
return MappingProxyType({})
def _loaded_catalog_entry(incoming_model_info: Mapping[str, object]) -> Mapping[str, object]:
LoadedCatalog: TypeAlias = Callable[[], Mapping[str, Mapping[str, object]]] # mutable-ok: Callable parameter syntax
def _loaded_catalog_entry(
incoming_model_info: Mapping[str, object], loaded_catalog: LoadedCatalog
) -> Mapping[str, object]:
catalog_key: Final = incoming_model_info.get(COST_MAP_LOOKUP_KEY)
if not isinstance(catalog_key, str):
return MappingProxyType({})
return GetModelCostMap.loaded_model_cost_map().get(catalog_key, MappingProxyType({}))
return loaded_catalog().get(catalog_key, MappingProxyType({}))
def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel:
def update_db_model(
db_model: Deployment,
updated_patch: updateDeployment,
loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map,
) -> PrismaCompatibleUpdateDBModel:
if updated_patch.model_info is not None:
_raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True))
merged_model_name: Final = updated_patch.model_name or db_model.model_name
@ -921,7 +930,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
echoed_fields: Final = echoed_cost_map_fields(
incoming_model_info,
_cost_map_entry(db_model, incoming_model_info),
_loaded_catalog_entry(incoming_model_info),
_loaded_catalog_entry(incoming_model_info, loaded_catalog),
)
merged_model_info.update(
MappingProxyType(

6
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -4119,15 +4119,12 @@ class TestModelInfoCostMapEchoFilter:
import litellm
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
bundled = litellm.get_model_info("openai/gpt-5.6")
remote = {**bundled, "max_input_tokens": bundled["max_input_tokens"] + 1}
monkeypatch.setattr(
GetModelCostMap, "_loaded_catalog", MappingProxyType({remote["key"]: MappingProxyType(remote)})
)
remote_catalog = MappingProxyType({remote["key"]: MappingProxyType(remote)})
monkeypatch.setattr(litellm, "get_model_info", lambda model, **_: {**remote, "max_input_tokens": 2048})
db_model = Deployment(
model_name="gpt-5.6",
@ -4138,6 +4135,7 @@ class TestModelInfoCostMapEchoFilter:
result = update_db_model(
db_model=db_model,
updated_patch=updateDeployment(model_info=ModelInfo(**{**remote, "id": "dep-echo-9", "db_model": True})),
loaded_catalog=lambda: remote_catalog,
)
info = json.loads(result["model_info"])