Merge pull request #5616 from BerriAI/litellm_fix_regen_keys

[Fix-Proxy] Regenerate keys when no duration is passed

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -293,9 +293,9 @@ async def prepare_key_update_data(
             continue
         if v is not None and v not in ([], {}, 0):
             non_default_values[k] = v
-
     if "duration" in non_default_values:
         duration = non_default_values.pop("duration")
+        if duration and (isinstance(duration, str)) and len(duration) > 0:
             duration_s = _duration_in_seconds(duration=duration)
             expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
             non_default_values["expires"] = expires

tests/proxy_admin_ui_tests/test_key_management.py

@@ -269,3 +269,56 @@ async def test_regenerate_api_key_with_new_alias_and_expiration(prisma_client):
     now = datetime.now(dt.timezone.utc)
     assert new_key.expires > now + dt.timedelta(days=29)
     assert new_key.expires < now + dt.timedelta(days=31)
+
+
+@pytest.mark.asyncio()
+async def test_regenerate_key_ui(prisma_client):
+    litellm.set_verbose = True
+    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()
+    import uuid
+
+    # generate new key
+    key_alias = f"test_alias_regenerate_key-{uuid.uuid4()}"
+    spend = 100
+    max_budget = 400
+    models = ["fake-openai-endpoint"]
+    new_key = await generate_key_fn(
+        data=GenerateKeyRequest(
+            key_alias=key_alias, spend=spend, max_budget=max_budget, models=models
+        ),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role=LitellmUserRoles.PROXY_ADMIN,
+            api_key="sk-1234",
+            user_id="1234",
+        ),
+    )
+
+    generated_key = new_key.key
+    print(generated_key)
+
+    # assert the new key works as expected
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    async def return_body():
+        return_string = f'{{"model": "fake-openai-endpoint"}}'
+        # return string as bytes
+        return return_string.encode()
+
+    request.body = return_body
+    result = await user_api_key_auth(request=request, api_key=f"Bearer {generated_key}")
+    print(result)
+
+    # regenerate the key
+    new_key = await regenerate_key_fn(
+        key=generated_key,
+        data=RegenerateKeyRequest(duration=""),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role=LitellmUserRoles.PROXY_ADMIN,
+            api_key="sk-1234",
+            user_id="1234",
+        ),
+    )
+    print("response from regenerate_key_fn", new_key)

tests/proxy_admin_ui_tests/test_usage_endpoints.py

@@ -89,6 +89,7 @@ from litellm.caching import DualCache
 from litellm.proxy._types import (
     DynamoDBArgs,
     GenerateKeyRequest,
+    RegenerateKeyRequest,
     KeyRequest,
     LiteLLM_UpperboundKeyGenerateParams,
     NewCustomerRequest,