From 2ad2d6cf9f7981e2dc1abe02b6efd748faf59658 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Apr 2024 08:17:45 -0700
Subject: [PATCH] return correct response from user/new

---
 litellm/proxy/proxy_server.py             |  7 +++++-
 litellm/tests/test_key_generate_prisma.py | 29 +++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index e580d3a8c03..7b693e73a9b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -5675,8 +5675,13 @@ async def new_user(data: NewUserRequest):
     return NewUserResponse(
         key=response.get("token", ""),
         expires=response.get("expires", None),
-        user_id=response["user_id"],
         max_budget=response["max_budget"],
+        user_id=response["user_id"],
+        team_id=response.get("team_id", None),
+        metadata=response.get("metadata", None),
+        models=response.get("models", None),
+        tpm_limit=response.get("tpm_limit", None),
+        rpm_limit=response.get("rpm_limit", None),
     )
 
 
diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py
index 683927ca92b..a50862bba46 100644
--- a/litellm/tests/test_key_generate_prisma.py
+++ b/litellm/tests/test_key_generate_prisma.py
@@ -108,6 +108,35 @@ def prisma_client():
     return prisma_client
 
 
+@pytest.mark.asyncio()
+async def test_new_user_response(prisma_client):
+    try:
+
+        print("prisma client=", prisma_client)
+
+        setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+        setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+
+        await litellm.proxy.proxy_server.prisma_client.connect()
+        from litellm.proxy.proxy_server import user_api_key_cache
+
+        _response = await new_user(
+            data=NewUserRequest(
+                models=["azure-gpt-3.5"],
+                team_id="ishaans-special-team",
+                tpm_limit=20,
+            )
+        )
+        print(_response)
+        assert _response.models == ["azure-gpt-3.5"]
+        assert _response.team_id == "ishaans-special-team"
+        assert _response.tpm_limit == 20
+
+    except Exception as e:
+        print("Got Exception", e)
+        pytest.fail(f"Got exception {e}")
+
+
 def test_generate_and_call_with_valid_key(prisma_client):
     # 1. Generate a Key, and use it to make a call