From fc16b6650e8e97ee2f4693a40f56c657c165b41c Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 6 Mar 2024 18:55:40 -0800
Subject: [PATCH] build(schema.prisma): add support for team-based model
 aliases

---
 litellm/proxy/schema.prisma | 15 ++++++++++++++-
 schema.prisma               | 15 ++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index 265bf32c0..031db99d1 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -42,6 +42,17 @@ model LiteLLM_OrganizationTable {
     teams LiteLLM_TeamTable[] 
 }
 
+// Model info for teams, just has model aliases for now.
+model LiteLLM_ModelTable {
+  id Int @id @default(autoincrement())
+  model_aliases Json? @map("aliases")
+  created_at    DateTime               @default(now()) @map("created_at")
+  created_by String
+  updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
+  updated_by String
+  team LiteLLM_TeamTable?
+}
+
 // Assign prod keys to groups, not individuals 
 model LiteLLM_TeamTable {
 		team_id    String @id @default(uuid())
@@ -63,7 +74,9 @@ model LiteLLM_TeamTable {
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
+    model_id Int? @unique
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
+    litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
 }
 
 // Track spend, rate limit, budget Users
@@ -149,4 +162,4 @@ model LiteLLM_UserNotifications {
   models              String[]
   justification       String
   status              String // approved, disapproved, pending
-}
+}
\ No newline at end of file
diff --git a/schema.prisma b/schema.prisma
index 265bf32c0..031db99d1 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -42,6 +42,17 @@ model LiteLLM_OrganizationTable {
     teams LiteLLM_TeamTable[] 
 }
 
+// Model info for teams, just has model aliases for now.
+model LiteLLM_ModelTable {
+  id Int @id @default(autoincrement())
+  model_aliases Json? @map("aliases")
+  created_at    DateTime               @default(now()) @map("created_at")
+  created_by String
+  updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
+  updated_by String
+  team LiteLLM_TeamTable?
+}
+
 // Assign prod keys to groups, not individuals 
 model LiteLLM_TeamTable {
 		team_id    String @id @default(uuid())
@@ -63,7 +74,9 @@ model LiteLLM_TeamTable {
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
+    model_id Int? @unique
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
+    litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
 }
 
 // Track spend, rate limit, budget Users
@@ -149,4 +162,4 @@ model LiteLLM_UserNotifications {
   models              String[]
   justification       String
   status              String // approved, disapproved, pending
-}
+}
\ No newline at end of file