From af2eabba91cdedb90df850704f170e4cd19cd383 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 29 Mar 2024 21:47:10 -0700 Subject: [PATCH] fix(proxy_server.py): fix `/key/update` endpoint to update key duration also adds a test for this to our ci/cd --- litellm/proxy/proxy_server.py | 7 +++++++ tests/test_keys.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4263675e6..5a3d72e51 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4189,6 +4189,13 @@ async def update_key_fn(request: Request, data: UpdateKeyRequest): 0, ): # models default to [], spend defaults to 0, we should not reset these values non_default_values[k] = v + + if "duration" in non_default_values: + duration = non_default_values.pop("duration") + duration_s = _duration_in_seconds(duration=duration) + expires = datetime.utcnow() + timedelta(seconds=duration_s) + non_default_values["expires"] = expires + response = await prisma_client.update_data( token=key, data={**non_default_values, "token": key} ) diff --git a/tests/test_keys.py b/tests/test_keys.py index 4a563cc18..33a5f2650 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -65,7 +65,7 @@ async def update_key(session, get_key): "Authorization": f"Bearer sk-1234", "Content-Type": "application/json", } - data = {"key": get_key, "models": ["gpt-4"]} + data = {"key": get_key, "models": ["gpt-4"], "duration": "120s"} async with session.post(url, headers=headers, json=data) as response: status = response.status