From 2b6972111e8b3256453b0acdec24c22232e16774 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 18 Jan 2024 14:42:46 -0800
Subject: [PATCH] (feat) write team_id to User Table

---
 litellm/proxy/proxy_server.py | 3 +++
 litellm/proxy/schema.prisma   | 1 +
 schema.prisma                 | 1 +
 3 files changed, 5 insertions(+)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index f95137630..80e89709d 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1081,12 +1081,15 @@ async def generate_key_helper_fn(
     config_json = json.dumps(config)
     metadata_json = json.dumps(metadata)
     user_id = user_id or str(uuid.uuid4())
+    if type(team_id) is not str:
+        team_id = str(team_id)
     try:
         # Create a new verification token (you may want to enhance this logic based on your needs)
         user_data = {
             "max_budget": max_budget,
             "user_email": user_email,
             "user_id": user_id,
+            "team_id": team_id,
             "spend": spend,
             "models": models,
         }
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index 1ed76140e..24f7f4f3d 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -9,6 +9,7 @@ generator client {
 
 model LiteLLM_UserTable {
 		user_id    String @unique
+    team_id    String?
 		max_budget Float?
     spend      Float    @default(0.0)
     user_email    String?
diff --git a/schema.prisma b/schema.prisma
index 3c2bf22cb..ae5c53a4d 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -9,6 +9,7 @@ generator client {
 
 model LiteLLM_UserTable {
 		user_id    String @unique
+    team_id    String?
 		max_budget Float?
     spend      Float    @default(0.0)
     user_email    String?