From 15685a8f53beaaf9015ae74ef4a449caf2dcc7b6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 2 Apr 2024 21:31:24 -0700 Subject: [PATCH] v0 use token_in /key_generate --- litellm/proxy/_types.py | 1 + litellm/proxy/proxy_server.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 757f89d4789..1486893ca6a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -306,6 +306,7 @@ class GenerateKeyResponse(GenerateKeyRequest): key_name: Optional[str] = None expires: Optional[datetime] user_id: Optional[str] = None + token_id: Optional[str] = None @root_validator(pre=True) def set_model_info(cls, values): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0923cfa7e75..4f4519035ee 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2563,7 +2563,10 @@ async def generate_key_helper_fn( ## CREATE KEY verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data) - await prisma_client.insert_data(data=key_data, table_name="key") + create_key_response = await prisma_client.insert_data( + data=key_data, table_name="key" + ) + key_data["token_id"] = create_key_response.token elif custom_db_client is not None: if table_name is None or table_name == "user": ## CREATE USER (If necessary)