diff --git a/litellm/models/end_user.py b/litellm/models/end_user.py index 349f1d8f2ce..3f711a31cfd 100644 --- a/litellm/models/end_user.py +++ b/litellm/models/end_user.py @@ -5,7 +5,7 @@ Canonical definition for ``litellm_endusertable``. Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, model_validator @@ -28,16 +28,18 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): @model_validator(mode="before") @classmethod - def set_model_info(cls, values): + def set_model_info(cls, values: Any) -> Any: if isinstance(values, BaseModel): - raw = values.model_dump() # rebind-ok: normalize model input + normalized = values.model_dump() elif hasattr(values, "__dict__") and not isinstance(values, dict): - raw = dict(values.__dict__) # rebind-ok: normalize object input # mutable-ok: object normalization + normalized = dict(values.__dict__) + elif isinstance(values, dict): + normalized = dict(values) else: - raw = values # rebind-ok: preserve input + return values - if isinstance(raw, dict) and raw.get("spend") is None: - raw.update({"spend": 0.0}) # mutable-ok: default spend value - return raw + if normalized.get("spend") is None: + normalized["spend"] = 0.0 + return normalized model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/models/tag.py b/litellm/models/tag.py index 701969e1999..b2da9f4642a 100644 --- a/litellm/models/tag.py +++ b/litellm/models/tag.py @@ -6,6 +6,7 @@ Canonical definition for ``litellm_tagtable``. Re-exported from """ from datetime import datetime +from typing import Any from pydantic import BaseModel, model_validator @@ -27,17 +28,18 @@ class LiteLLM_TagTable(LiteLLMPydanticObjectBase): @model_validator(mode="before") @classmethod - def set_model_info(cls, values): + def set_model_info(cls, values: Any) -> Any: if isinstance(values, BaseModel): - raw = values.model_dump() # rebind-ok: normalize model input + normalized = values.model_dump() elif hasattr(values, "__dict__") and not isinstance(values, dict): - raw = dict(values.__dict__) # rebind-ok: normalize object input # mutable-ok: object normalization + normalized = dict(values.__dict__) + elif isinstance(values, dict): + normalized = dict(values) else: - raw = values # rebind-ok: preserve input + return values - if isinstance(raw, dict): - if raw.get("spend") is None: - raw.update({"spend": 0.0}) # mutable-ok: default spend value - if raw.get("models") is None: - raw.update({"models": []}) # mutable-ok: default models value - return raw + if normalized.get("spend") is None: + normalized["spend"] = 0.0 + if normalized.get("models") is None: + normalized["models"] = [] + return normalized diff --git a/litellm/models/user.py b/litellm/models/user.py index 8c3bf065c12..e2aeadda793 100644 --- a/litellm/models/user.py +++ b/litellm/models/user.py @@ -6,6 +6,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from """ from datetime import datetime +from typing import Any from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -49,22 +50,23 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): @model_validator(mode="before") @classmethod - def set_model_info(cls, values): + def set_model_info(cls, values: Any) -> Any: if isinstance(values, BaseModel): - raw = values.model_dump() # rebind-ok: normalize model input + normalized = values.model_dump() elif hasattr(values, "__dict__") and not isinstance(values, dict): - raw = dict(values.__dict__) # rebind-ok: normalize object input # mutable-ok: object normalization + normalized = dict(values.__dict__) + elif isinstance(values, dict): + normalized = dict(values) else: - raw = values # rebind-ok: preserve input + return values - if isinstance(raw, dict): - if raw.get("spend") is None: - raw.update({"spend": 0.0}) # mutable-ok: default spend value - if raw.get("models") is None: - raw.update({"models": []}) # mutable-ok: default models value - if raw.get("teams") is None: - raw.update({"teams": []}) # mutable-ok: default teams value - return raw + if normalized.get("spend") is None: + normalized["spend"] = 0.0 + if normalized.get("models") is None: + normalized["models"] = [] + if normalized.get("teams") is None: + normalized["teams"] = [] + return normalized def is_over_budget(self) -> bool: if self.max_budget is None: diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index d3a569b8691..8ff8061164b 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -307,6 +307,11 @@ class TestUser: def test_user_from_object_with_attributes(self): class MockPrismaUser: + user_id: str + spend: float | None + models: list[str] | None + teams: list[str] | None + def __init__( self, user_id: str, @@ -342,6 +347,10 @@ class TestUser: assert user.models == [] assert user.teams == [] + def test_user_from_invalid_type_raises_validation_error(self): + with pytest.raises(ValidationError): + LiteLLM_UserTable.model_validate(123) + def test_password_hash_excluded_from_serialization(self): from litellm.proxy._types import LiteLLM_UserTableWithKeyCount @@ -457,6 +466,10 @@ class TestTagTable: def test_tag_from_object_with_attributes(self): class MockPrismaTag: + tag_name: str + spend: float | None + models: list[str] | None + def __init__( self, tag_name: str, @@ -485,6 +498,10 @@ class TestTagTable: assert tag.spend == 0.0 assert tag.models == [] + def test_tag_from_invalid_type_raises_validation_error(self): + with pytest.raises(ValidationError): + LiteLLM_TagTable.model_validate(123) + class TestEndUserTable: def test_end_user_creation(self): @@ -506,6 +523,11 @@ class TestEndUserTable: def test_end_user_from_object_with_attributes(self): class MockPrismaEndUser: + user_id: str + blocked: bool + spend: float | None + alias: str | None + def __init__( self, user_id: str, @@ -536,6 +558,10 @@ class TestEndUserTable: assert eu.blocked is False assert eu.spend == 0.0 + def test_end_user_from_invalid_type_raises_validation_error(self): + with pytest.raises(ValidationError): + LiteLLM_EndUserTable.model_validate(123) + class TestBudgetTableFull: def test_full_adds_server_managed_fields(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 2df4237bda5..66d77d9faab 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -473,11 +473,22 @@ def test_block_customer_success_with_prisma_row_object(mock_prisma_client, mock_ BlockUsersResponse must serialize them without raising 500. """ class MockPrismaRow: - def __init__(self, user_id: str, blocked: bool, spend: float | None = None) -> None: - self.user_id = user_id - self.blocked = blocked - self.spend = spend - self.alias = None + user_id: str + blocked: bool + spend: float | None + alias: str | None + + def __init__( + self, + user_id: str, + blocked: bool, + spend: float | None = None, + alias: str | None = None, + ) -> None: + self.user_id: str = user_id + self.blocked: bool = blocked + self.spend: float | None = spend + self.alias: str | None = alias blocked_row = MockPrismaRow(user_id="blocked-prisma-1", blocked=True, spend=None) mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(return_value=blocked_row)