feat(litellm): finish models/repository migration; route all table access through repositories - #29713
Closed
yassin-berriai wants to merge 35 commits into
Closed
feat(litellm): finish models/repository migration; route all table access through repositories#29713yassin-berriai wants to merge 35 commits into
yassin-berriai wants to merge 35 commits into
Conversation
This commit introduces domain models under litellm/backend and repository layer under litellm/gateway to centralize database operations and provide a clean separation between business logic and data access. Backend models (litellm/backend/models/): - Budget: budget configuration with spend tracking - Credentials: encrypted credential storage - Model: proxy model definitions with encryption support - ObjectPermission: MCP, vector store, and agent permissions - Organization: organization management - Project: project management between teams and keys - Team: team management with member roles - User: user management with budget tracking - VerificationToken: API key/token management Gateway repositories (litellm/gateway/repositories/): - BaseRepository: abstract base with common CRUD operations - BudgetRepository: budget table operations - ConfigRepository: config reconciliation from DB and configmap - CredentialsRepository: encrypted credential operations - ModelRepository: proxy model operations with encryption - ObjectPermissionRepository: permission table operations - OrganizationRepository: organization operations - ProjectRepository: project operations - TeamRepository: team operations with audit logging - UserRepository: user operations - VerificationTokenRepository: token operations with audit logging The ConfigRepository implements the config reconciliation strategy where DB values override YAML configmap values, except for None values and empty lists which preserve the YAML config. Includes comprehensive unit tests for all models and repositories. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Replace read-modify-write patterns with atomic array push operations for add_member, add_admin, add_models in TeamRepository and add_to_team in UserRepository. Use transactions for delete_team and delete_token to ensure archive-then-delete is atomic. Add scalability note for find_by_team_id in ModelRepository. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Add type assertions for _to_model returns in create methods, fix _to_model_list to properly filter None values, and cast param_name to Literal type in config reconciliation. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Include tests/test_litellm/backend and tests/test_litellm/gateway in the unit test misc workflow so the new repository and model tests are picked up by CI coverage. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Add comprehensive tests for all repository types including organization, project, object_permission, and credentials. Fix credentials_repository _to_model to properly parse JSON fields. Tests increased from 46 to 115. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Add missing organization_id and budget_id fields to VerificationToken model. The DB columns were being silently dropped due to extra="ignore" in the base model. Keep org_id for API compatibility while organization_id matches the actual DB column name. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Both backend (management endpoints) and gateway routes use models and repositories, so they should not be nested under either. This moves litellm/backend/models/ to litellm/models/ and litellm/gateway/repositories/ to litellm/repositories/. The backend and gateway __init__.py files remain as backwards-compatible re-exports from the new locations. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Tests for ConfigRepository reconcile_config, update_config_fields, and environment variable handling. Tests for VerificationTokenRepository and TeamRepository delete operations with transactions, JSON field parsing, and extended update operations. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Tests for ModelRepository.find_by_team_id, BaseRepository pagination and filtering, DomainModel.from_db_record edge cases. Coverage should now exceed the threshold. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Tests all field branches in update_team, update_project, update_organization, and update_permission to improve patch coverage. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Tests create_model with all fields, update_model with all fields, and update_user with all fields to cover remaining branch coverage gaps. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
…llow copy Two bugs fixed based on Greptile review: 1. delete_team was using to_db_dict() which includes fields like default_team_member_models and budget_limits that don't exist in LiteLLM_DeletedTeamTable. Added _build_archive_data() method that explicitly includes only columns present in the archive table. 2. reconcile_config used yaml_config.copy() (shallow copy) which caused _deep_merge_dicts to silently mutate the caller's nested dicts. Changed to copy.deepcopy() to ensure the original config is not modified. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
These were added as backwards-compatible re-exports but nothing uses them. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Move the canonical credential types (CredentialBase, CredentialItem, CreateCredentialItem) into the model layer (litellm/models/credentials.py) and re-export them from litellm.types.utils for backwards compatibility. Replace the hand-rolled Credentials domain model with these. CredentialsRepository is now the only code path that touches litellm_credentialstable; the credential CRUD endpoints and proxy_server.get_credentials go through it. Encryption and in-memory credential_list syncing stay with the callers so stored-data behavior is unchanged. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Standardize on LiteLLM_ObjectPermissionTable as the canonical type, now defined in litellm/models/object_permission.py (subclassing the low-level LiteLLMPydanticObjectBase) and re-exported from litellm.proxy._types. This also reconciles a latent bug: the table has a models column that the old _types definition omitted. Relax BaseRepository's generic bound from DomainModel to BaseModel and make the default record->model conversion generic so repositories can return the canonical table types. Repoint ObjectPermissionRepository accordingly. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Standardize on LiteLLM_BudgetTable (the user-input allowlist base) as the canonical type, now defined in litellm/models/budget.py and re-exported from litellm.proxy._types; its API-request subclasses stay in _types. Repoint BudgetRepository. Drop the hand-rolled Budget domain model and its unused helper methods. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Standardize on LiteLLM_ProxyModelTable, now defined in litellm/models/model.py and re-exported from litellm.proxy._types. Reconciles latent omissions vs the schema (blocked column, nullable model_info) and keeps the team_id / team_public_model_name / is_blocked helpers the repository relies on. Repoint ModelRepository; drop the hand-rolled Model domain model. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Standardize on LiteLLM_ProjectTable, now defined in litellm/models/project.py (referencing the relocated Budget and ObjectPermission models) and re-exported from litellm.proxy._types. Repoint ProjectRepository; drop the hand-rolled Project domain model. Teach the test MockTable to synthesize primary keys on create so repositories returning PK-required canonical types are covered. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
model_id is the primary key and always present on reads; restoring it to a required str matches the original _types contract and fixes the audit-log call site type error. The test MockTable synthesizes model_id on create. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
…layer Standardize on LiteLLM_UserTable, now defined in litellm/models/user.py and re-exported from litellm.proxy._types. Merge in the schema columns the prior canonical class omitted (team_id, organization_id, object_permission_id, password, max_parallel_requests, allowed_cache_controls, policies) so reads through UserRepository no longer drop data. Relocate the referenced LiteLLM_OrganizationMembershipTable alongside it. Repoint UserRepository; drop the hand-rolled User domain model. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Standardize on LiteLLM_OrganizationTable, now defined in litellm/models/organization.py (referencing the relocated Budget, ObjectPermission and User models) and re-exported from litellm.proxy._types. Reconciles the missing model_spend column. Repoint OrganizationRepository; drop the hand-rolled Organization domain model. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Standardize on LiteLLM_VerificationToken (and LiteLLM_DeletedVerificationToken), now defined in litellm/models/verification_token.py and re-exported from litellm.proxy._types; LiteLLM_VerificationTokenView stays in _types since it references Member. Repoint VerificationTokenRepository and build the archive record via model_dump (dropping the object_permission relation) instead of the removed DomainModel helper. The canonical class is moved verbatim, preserving its existing org_id semantics rather than changing auth-facing behavior. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
…model layer Move the team cluster (MemberBase, Member, BudgetLimitEntry, LiteLLM_ModelTable, TeamBase, LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_DeletedTeamTable) into litellm/models/team.py, re-exported from litellm.proxy._types so the many importers and request-model subclasses (OrgMember, NewTeamRequest, etc.) keep working. Reconciles columns the canonical class omitted but the repository's archive path uses (model_spend, model_max_budget, policies, allow_team_guardrail_config). Repoint TeamRepository; drop the hand-rolled Team domain model. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
The Team relocation script's class-replacement consumed the following LiteLLM_ProxyModelTable re-export line, breaking import of litellm.proxy.proxy_server across the proxy test suites. Restore it. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
Relocate LiteLLM_Config, LiteLLM_SkillsTable, LiteLLM_AccessGroupTable, LiteLLM_TagTable and LiteLLM_EndUserTable into litellm/models/, re-exported from litellm.proxy._types. Also reformat _types.py. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
policies was reconciled as a required List, but team creation passes policies=None; make it Optional[List[str]] to restore team creation. Harden the team archive build to coerce None arrays/spend to schema-safe defaults. Add unit tests for the relocated Config, Skills, AccessGroup, Tag and EndUser table models. https://claude.ai/code/session_01MtrRXaRiaAgKqF4UcbyKhC
…le tables into the model layer Moves the remaining standalone DB-table pydantic definitions out of litellm/proxy/_types.py and into litellm/models/, re-exported from _types for backwards compatibility. Also moves LiteLLM_BudgetTableFull and LiteLLM_TeamMemberTable next to LiteLLM_BudgetTable, and completes the models package __init__ exports. Imports MCPAuthType/MCPTransportType in mcp_server_manager from their source litellm.types.mcp instead of the proxy._types re-export, breaking an import cycle that otherwise surfaced once the model layer pulled in mcp_server.
…pository Replaces direct prisma_client.db.litellm_objectpermissiontable access across the proxy with ObjectPermissionRepository(...).table, centralizing the table name in the repository as the single data-access chokepoint. Behavior is unchanged; the property returns the same Prisma delegate.
Replaces direct prisma_client.db.litellm_<table> access across the proxy with <Table>Repository(...).table so each table is named in exactly one place, the repository layer. For tables that already had a domain repository (team, user, verification token, budget, model, organization, project, config) the existing class gains the call sites; the rest get a thin PrismaTableRepository subclass in table_repositories.py. The .table property returns the same Prisma delegate, so behavior is unchanged. Adds tests covering the passthrough base (table binding, no-DB guard, and that every generated repository binds a distinct litellm_ table name).
|
|
The generated passthrough repository module was committed without black formatting (it was untracked when the earlier black pass ran over the tracked diff), so it lacked the blank lines black wants between top-level classes.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Relevant issues
Continues the models and repository layer effort from #29686
Linear ticket
https://linear.app/litellm-ai/issue/LIT-3570/create-litellm-models-and-repository-layers
Pre-Submission checklist
make test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewType
New Feature
Changes
This picks up where #29686 left off and finishes the migration so that every DB table definition lives in
litellm/models/and every inline Prisma table access goes through the repository layer.Definitions: the remaining pydantic table classes move out of
litellm/proxy/_types.pyintolitellm/models/and are re-exported for backwards compatibility. That covers the MCP server table, spend and error logs, team membership, the managed file/object/vector-store tables, andLiteLLM_BudgetTableFullplusLiteLLM_TeamMemberTablewhich now sit next to the budget model. Themodelspackage__init__is completed to export every domain model. While wiringmcp_serverinto the model layer surfaced an import cycle, somcp_server_managernow importsMCPAuthType/MCPTransportTypefrom their sourcelitellm.types.mcprather than theproxy._typesre-export.Operations: every
prisma_client.db.litellm_<table>.<op>(...)call across the proxy now goes through<Table>Repository(...).table, so each table name is defined in exactly one place. The nine tables that already had a domain repository (team, user, verification token, budget, model, organization, project, config, object permission) get their call sites routed in; the remaining tables get a thinPrismaTableRepositorysubclass inrepositories/table_repositories.py. The.tableproperty returns the same Prisma delegate, so this is behavior preserving.Two categories are intentionally left as direct access because they do not map to the non-transactional
.tabledelegate: transaction and batch operations insidebatch_()/tx()blocks (batcher.litellm_X,tx.litellm_X), and raw SQL viaquery_raw/execute_rawthat names tables as SQL strings. Routing those through.tablewould break transactional semantics or would not apply.Tests cover the relocated models (including the team-membership budget helpers and the managed tables) and the passthrough repository base (table binding, the no-DB guard, and that every generated repository binds a distinct
litellm_table name).Screenshots / Proof of Fix
To be added with a live proxy run against real provider APIs
Generated by Claude Code