From 8c911568421343061366336f8e45a2bed39c22ef Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 19 Mar 2024 16:59:59 -0700
Subject: [PATCH] build: build fixes

---
 Dockerfile                       |  5 +++++
 litellm/proxy/auth/handle_jwt.py |  2 +-
 litellm/proxy/utils.py           | 18 ++++++++++++------
 requirements.txt                 |  1 -
 4 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 4408afb3d..f560473e2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -38,6 +38,11 @@ RUN pip wheel --no-cache-dir --wheel-dir=/wheels/ -r requirements.txt
 # install semantic-cache [Experimental]- we need this here and not in requirements.txt because redisvl pins to pydantic 1.0 
 RUN pip install redisvl==0.0.7 --no-deps
 
+# ensure pyjwt is used, not jwt
+RUN pip uninstall jwt -y
+RUN pip uninstall PyJWT -y
+RUN pip install PyJWT --no-cache-dir
+
 # Build Admin UI
 RUN chmod +x build_admin_ui.sh && ./build_admin_ui.sh
 
diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index 108231df6..4c3c013c7 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -8,8 +8,8 @@ JWT token must have 'litellm_proxy_admin' in scope.
 
 import httpx
 import jwt
-import json
 from jwt.algorithms import RSAAlgorithm
+import json
 import os
 from litellm.proxy._types import LiteLLMProxyRoles
 from typing import Optional
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 879caa495..32289cb2f 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -1870,7 +1870,9 @@ async def update_spend(
     if len(prisma_client.user_list_transactons.keys()) > 0:
         for i in range(n_retry_times + 1):
             try:
-                async with prisma_client.db.tx(timeout=6000) as transaction:
+                async with prisma_client.db.tx(
+                    timeout=timedelta(seconds=60)
+                ) as transaction:
                     async with transaction.batch_() as batcher:
                         for (
                             user_id,
@@ -1895,7 +1897,9 @@ async def update_spend(
     if len(prisma_client.end_user_list_transactons.keys()) > 0:
         for i in range(n_retry_times + 1):
             try:
-                async with prisma_client.db.tx(timeout=6000) as transaction:
+                async with prisma_client.db.tx(
+                    timeout=timedelta(seconds=60)
+                ) as transaction:
                     async with transaction.batch_() as batcher:
                         for (
                             end_user_id,
@@ -1926,7 +1930,9 @@ async def update_spend(
     if len(prisma_client.key_list_transactons.keys()) > 0:
         for i in range(n_retry_times + 1):
             try:
-                async with prisma_client.db.tx(timeout=6000) as transaction:
+                async with prisma_client.db.tx(
+                    timeout=timedelta(seconds=60)
+                ) as transaction:
                     async with transaction.batch_() as batcher:
                         for (
                             token,
@@ -1951,7 +1957,9 @@ async def update_spend(
     if len(prisma_client.team_list_transactons.keys()) > 0:
         for i in range(n_retry_times + 1):
             try:
-                async with prisma_client.db.tx(timeout=6000) as transaction:
+                async with prisma_client.db.tx(
+                    timeout=timedelta(seconds=60)
+                ) as transaction:
                     async with transaction.batch_() as batcher:
                         for (
                             team_id,
@@ -1972,8 +1980,6 @@ async def update_spend(
             except Exception as e:
                 raise e
 
-    ### UPDATE SPEND LOGS TABLE ###
-
 
 async def monitor_spend_list(prisma_client: PrismaClient):
     """
diff --git a/requirements.txt b/requirements.txt
index 47889682f..5d6b2307d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,7 +24,6 @@ apscheduler==3.10.4 # for resetting budget in background
 fastapi-sso==0.10.0 # admin UI, SSO
 PyJWT==2.8.0 # admin UI, jwts 
 python-multipart==0.0.6 # admin UI
-jwt==1.3.1 # for jwt auth
 ### LITELLM PACKAGE DEPENDENCIES
 python-dotenv>=0.2.0 # for env 
 tiktoken>=0.4.0 # for calculating usage