feat(litellm): finish models/repository migration; route all table access through repositories

.github/workflows/test-unit-misc.yml

@@ -28,6 +28,8 @@ jobs:
         tests/test_litellm/completion_extras
         tests/test_litellm/containers
         tests/test_litellm/experimental_mcp_client
+        tests/test_litellm/models
+        tests/test_litellm/repositories
         tests/test_litellm/images
         tests/test_litellm/interactions
         tests/test_litellm/passthrough

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -37,6 +37,8 @@
     VirtualKeyEvent,
     WebhookEvent,
 )
+from litellm.repositories.team_repository import TeamRepository
+from litellm.repositories.user_repository import UserRepository
 from litellm.types.integrations.slack_alerting import *
 
 from ..email_templates.templates import *
@@ -1231,7 +1233,7 @@ async def send_key_created_or_user_invited_email(
                 and recipient_user_id is not None
                 and prisma_client is not None
             ):
-                user_row = await prisma_client.db.litellm_usertable.find_unique(
+                user_row = await UserRepository(prisma_client).table.find_unique(
                     where={"user_id": recipient_user_id}
                 )
 
@@ -1263,7 +1265,7 @@ async def send_key_created_or_user_invited_email(
                 team_id = webhook_event.team_id
                 team_name = "Default Team"
                 if team_id is not None and prisma_client is not None:
-                    team_row = await prisma_client.db.litellm_teamtable.find_unique(
+                    team_row = await TeamRepository(prisma_client).table.find_unique(
                         where={"team_id": team_id}
                     )
                     if team_row is not None:

litellm/integrations/email_alerting.py

@@ -7,6 +7,7 @@
 
 from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm.proxy._types import WebhookEvent
+from litellm.repositories.team_repository import TeamRepository
 
 # we use this for the email header, please send a test email if you change this. verify it looks good on email
 LITELLM_LOGO_URL = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
@@ -24,7 +25,7 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list:
     if prisma_client is None:
         raise Exception("Not connected to DB!")
 
-    team_row = await prisma_client.db.litellm_teamtable.find_unique(
+    team_row = await TeamRepository(prisma_client).table.find_unique(
         where={
             "team_id": team_id,
         }

litellm/integrations/prometheus.py

@@ -25,13 +25,13 @@
 import litellm
 from litellm._logging import print_verbose, verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
-    BoundedPrometheusSeriesTracker,
-)
 from litellm.integrations.prometheus_helpers import (
     PrometheusLabelFactoryContext,
     _get_cached_end_user_id_for_cost_tracking,
 )
+from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
+    BoundedPrometheusSeriesTracker,
+)
 from litellm.litellm_core_utils.core_helpers import (
     get_litellm_metadata_from_kwargs,
     get_metadata_variable_name_from_kwargs,
@@ -42,6 +42,9 @@
     LiteLLM_UserTable,
     UserAPIKeyAuth,
 )
+from litellm.repositories.organization_repository import OrganizationRepository
+from litellm.repositories.team_repository import TeamRepository
+from litellm.repositories.user_repository import UserRepository
 from litellm.types.integrations.prometheus import *
 from litellm.types.integrations.prometheus import (
     _sanitize_prometheus_label_name,
@@ -3198,12 +3201,12 @@ async def fetch_users(
             page_size: int, page: int
         ) -> Tuple[List[LiteLLM_UserTable], Optional[int]]:
             skip = (page - 1) * page_size
-            users = await prisma_client.db.litellm_usertable.find_many(
+            users = await UserRepository(prisma_client).table.find_many(
                 skip=skip,
                 take=page_size,
                 order={"created_at": "desc"},
             )
-            total_count = await prisma_client.db.litellm_usertable.count()
+            total_count = await UserRepository(prisma_client).table.count()
             return users, total_count
 
         await self._initialize_budget_metrics(
@@ -3226,13 +3229,13 @@ async def _initialize_org_budget_metrics(self):
 
         async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]:
             skip = (page - 1) * page_size
-            orgs = await prisma_client.db.litellm_organizationtable.find_many(
+            orgs = await OrganizationRepository(prisma_client).table.find_many(
                 skip=skip,
                 take=page_size,
                 order={"created_at": "desc"},
                 include={"litellm_budget_table": True},
             )
-            total_count = await prisma_client.db.litellm_organizationtable.count()
+            total_count = await OrganizationRepository(prisma_client).table.count()
             return orgs, total_count
 
         await self._initialize_budget_metrics(
@@ -3300,14 +3303,14 @@ async def _initialize_user_and_team_count_metrics(self):
 
         try:
             # Get total user count
-            total_users = await prisma_client.db.litellm_usertable.count()
+            total_users = await UserRepository(prisma_client).table.count()
             self.litellm_total_users_metric.set(total_users)
             verbose_logger.debug(
                 f"Prometheus: set litellm_total_users to {total_users}"
             )
 
             # Get total team count
-            total_teams = await prisma_client.db.litellm_teamtable.count()
+            total_teams = await TeamRepository(prisma_client).table.count()
             self.litellm_teams_count_metric.set(total_teams)
             verbose_logger.debug(
                 f"Prometheus: set litellm_teams_count to {total_teams}"

litellm/llms/litellm_proxy/skills/handler.py

@@ -17,6 +17,7 @@
     is_proxy_admin,
     user_can_access_resource_owner,
 )
+from litellm.repositories.table_repositories import SkillsRepository
 
 # Skills are looked up on every chat completion that has skills enabled
 # (`SkillsInjectionHook` calls ``fetch_skill_from_db``). 60s LRU/TTL cache
@@ -107,7 +108,7 @@ async def create_skill(
             f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}"
         )
 
-        new_skill = await prisma_client.db.litellm_skillstable.create(data=skill_data)
+        new_skill = await SkillsRepository(prisma_client).table.create(data=skill_data)
         return _prisma_skill_to_litellm(new_skill)
 
     @staticmethod
@@ -133,7 +134,7 @@ async def list_skills(
                 return []
             find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}}
 
-        skills = await prisma_client.db.litellm_skillstable.find_many(
+        skills = await SkillsRepository(prisma_client).table.find_many(
             **find_many_kwargs
         )
         return [_prisma_skill_to_litellm(s) for s in skills]
@@ -150,7 +151,7 @@ async def _load_skill(skill_id: str) -> Optional[Any]:
             return cached
 
         prisma_client = await LiteLLMSkillsHandler._get_prisma_client()
-        skill = await prisma_client.db.litellm_skillstable.find_unique(
+        skill = await SkillsRepository(prisma_client).table.find_unique(
             where={"skill_id": skill_id}
         )
         _SKILL_CACHE.set_cache(
@@ -189,7 +190,7 @@ async def delete_skill(
         ):
             raise ValueError(f"Skill not found: {skill_id}")
 
-        await prisma_client.db.litellm_skillstable.delete(where={"skill_id": skill_id})
+        await SkillsRepository(prisma_client).table.delete(where={"skill_id": skill_id})
         _SKILL_CACHE.set_cache(skill_id, _NEGATIVE_SKILL_SENTINEL)
 
         return {"id": skill_id, "type": "skill_deleted"}

litellm/models/__init__.py

@@ -0,0 +1,66 @@
+"""
+Domain models for LiteLLM backend.
+"""
+
+from litellm.models.access_group import LiteLLM_AccessGroupTable
+from litellm.models.budget import (
+    LiteLLM_BudgetTable,
+    LiteLLM_BudgetTableFull,
+    LiteLLM_TeamMemberTable,
+)
+from litellm.models.config import LiteLLM_Config
+from litellm.models.credentials import (
+    CreateCredentialItem,
+    CredentialBase,
+    CredentialItem,
+)
+from litellm.models.end_user import LiteLLM_EndUserTable
+from litellm.models.managed_files import (
+    LiteLLM_ManagedFileTable,
+    LiteLLM_ManagedObjectTable,
+    LiteLLM_ManagedVectorStoresTable,
+    LiteLLM_ManagedVectorStoreTable,
+)
+from litellm.models.mcp_server import LiteLLM_MCPServerTable
+from litellm.models.model import LiteLLM_ProxyModelTable
+from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
+from litellm.models.organization import LiteLLM_OrganizationTable
+from litellm.models.organization_membership import LiteLLM_OrganizationMembershipTable
+from litellm.models.project import LiteLLM_ProjectTable
+from litellm.models.skills import LiteLLM_SkillsTable
+from litellm.models.spend_logs import LiteLLM_ErrorLogs, LiteLLM_SpendLogs
+from litellm.models.tag import LiteLLM_TagTable
+from litellm.models.team import LiteLLM_TeamTable
+from litellm.models.team_membership import LiteLLM_TeamMembership
+from litellm.models.user import LiteLLM_UserTable
+from litellm.models.verification_token import LiteLLM_VerificationToken
+
+__all__ = [
+    "LiteLLM_AccessGroupTable",
+    "LiteLLM_BudgetTable",
+    "LiteLLM_BudgetTableFull",
+    "LiteLLM_TeamMemberTable",
+    "LiteLLM_Config",
+    "CredentialBase",
+    "CredentialItem",
+    "CreateCredentialItem",
+    "LiteLLM_EndUserTable",
+    "LiteLLM_ManagedFileTable",
+    "LiteLLM_ManagedObjectTable",
+    "LiteLLM_ManagedVectorStoreTable",
+    "LiteLLM_ManagedVectorStoresTable",
+    "LiteLLM_MCPServerTable",
+    "LiteLLM_ProxyModelTable",
+    "LiteLLM_ObjectPermissionTable",
+    "LiteLLM_OrganizationTable",
+    "LiteLLM_OrganizationMembershipTable",
+    "LiteLLM_ProjectTable",
+    "LiteLLM_SkillsTable",
+    "LiteLLM_ErrorLogs",
+    "LiteLLM_SpendLogs",
+    "LiteLLM_TagTable",
+    "LiteLLM_TeamTable",
+    "LiteLLM_TeamMembership",
+    "LiteLLM_UserTable",
+    "LiteLLM_VerificationToken",
+]

litellm/models/access_group.py

@@ -0,0 +1,26 @@
+"""
+Access group table model.
+
+Canonical definition for ``litellm_accessgrouptable``. Re-exported from
+``litellm.proxy._types`` for backwards compatibility.
+"""
+
+from datetime import datetime
+from typing import List, Optional
+
+from litellm.types.llms.base import LiteLLMPydanticObjectBase
+
+
+class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase):
+    access_group_id: str
+    access_group_name: str
+    description: Optional[str] = None
+    access_model_names: List[str] = []
+    access_mcp_server_ids: List[str] = []
+    access_agent_ids: List[str] = []
+    assigned_team_ids: List[str] = []
+    assigned_key_ids: List[str] = []
+    created_at: Optional[datetime] = None
+    created_by: Optional[str] = None
+    updated_at: Optional[datetime] = None
+    updated_by: Optional[str] = None

litellm/models/base.py

@@ -0,0 +1,38 @@
+"""
+Base model class for domain models.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+
+class DomainModel(BaseModel):
+    """Base class for all domain models."""
+
+    model_config = ConfigDict(
+        from_attributes=True,
+        protected_namespaces=(),
+        extra="ignore",
+    )
+
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    @classmethod
+    def from_db_record(cls, record: Any) -> "DomainModel":
+        """Create a domain model from a database record."""
+        if record is None:
+            raise ValueError("Cannot create domain model from None record")
+        if isinstance(record, dict):
+            return cls(**record)
+        if hasattr(record, "model_dump") and callable(record.model_dump):
+            return cls(**record.model_dump())
+        if hasattr(record, "dict") and callable(record.dict):
+            return cls(**record.dict())
+        return cls(**dict(record))
+
+    def to_db_dict(self, exclude_unset: bool = False) -> Dict[str, Any]:
+        """Convert domain model to a dictionary for database operations."""
+        return self.model_dump(exclude_none=True, exclude_unset=exclude_unset)

litellm/models/budget.py

@@ -0,0 +1,56 @@
+"""
+Budget table model.
+
+Canonical definition for ``litellm_budgettable``. Re-exported from
+``litellm.proxy._types`` for backwards compatibility.
+"""
+
+from datetime import datetime
+from typing import List, Optional
+
+from pydantic import ConfigDict
+
+from litellm.types.llms.base import LiteLLMPydanticObjectBase
+
+
+class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
+    """Represents user-controllable params for a LiteLLM_BudgetTable record.
+
+    Budget-write paths use `model_fields.keys()` on this class as an allowlist
+    for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
+    `LiteLLM_BudgetTableFull` so they aren't user-settable.
+    """
+
+    budget_id: Optional[str] = None
+    soft_budget: Optional[float] = None
+    max_budget: Optional[float] = None
+    max_parallel_requests: Optional[int] = None
+    tpm_limit: Optional[int] = None
+    rpm_limit: Optional[int] = None
+    model_max_budget: Optional[dict] = None
+    budget_duration: Optional[str] = None
+    allowed_models: Optional[List[str]] = (
+        None  # per-member model scope; empty = inherit team models
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
+    """LiteLLM_BudgetTable + server-managed fields returned on API responses."""
+
+    budget_reset_at: Optional[datetime] = None
+    created_at: datetime
+
+
+class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable):
+    """
+    Used to track spend of a user_id within a team_id
+    """
+
+    spend: Optional[float] = None
+    user_id: Optional[str] = None
+    team_id: Optional[str] = None
+    budget_id: Optional[str] = None
+
+    model_config = ConfigDict(protected_namespaces=())

litellm/models/config.py

@@ -0,0 +1,15 @@
+"""
+Config table model.
+
+Canonical definition for ``litellm_config``. Re-exported from
+``litellm.proxy._types`` for backwards compatibility.
+"""
+
+from typing import Dict
+
+from litellm.types.llms.base import LiteLLMPydanticObjectBase
+
+
+class LiteLLM_Config(LiteLLMPydanticObjectBase):
+    param_name: str
+    param_value: Dict