Make key creation idempotent (reuse existing key by alias)
This commit is contained in:
1 parent
f3cc8e1d89
commit
6257c89800
1 file changed
+35
-2
+35
-2
@@ -222,14 +222,47 @@ async def cloudron_send_invite(user_id: str, email: str) -> None:
|
||||
r.raise_for_status()
|
||||
|
||||
|
||||
async def litellm_find_key_by_alias(alias: str) -> str | None:
|
||||
"""Find an existing key's token by its alias (key/list returns tokens)."""
|
||||
async with httpx.AsyncClient() as client:
|
||||
r = await client.get(
|
||||
f"{LITELLM_BASE}/key/list",
|
||||
headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"},
|
||||
)
|
||||
r.raise_for_status()
|
||||
tokens = r.json().get("keys", [])
|
||||
for token in tokens:
|
||||
info = await client.get(
|
||||
f"{LITELLM_BASE}/key/info",
|
||||
headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"},
|
||||
params={"key": token},
|
||||
)
|
||||
if info.status_code == 200:
|
||||
data = info.json().get("info", {})
|
||||
if data.get("key_alias") == alias:
|
||||
return token
|
||||
return None
|
||||
|
||||
|
||||
async def litellm_create_key(email: str, budget: float) -> str:
|
||||
"""Create a LiteLLM virtual key for a member with a budget cap."""
|
||||
"""Create a LiteLLM virtual key for a member with a budget cap.
|
||||
|
||||
Idempotent: if a key with this alias already exists, return its token.
|
||||
"""
|
||||
alias = f"member:{email}"
|
||||
|
||||
# If the key already exists (e.g. from a prior partial attempt), reuse it.
|
||||
existing = await litellm_find_key_by_alias(alias)
|
||||
if existing:
|
||||
logger.info("Reusing existing LiteLLM key for %s", email)
|
||||
return existing
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
r = await client.post(
|
||||
f"{LITELLM_BASE}/key/generate",
|
||||
headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"},
|
||||
json={
|
||||
"key_alias": f"member:{email}",
|
||||
"key_alias": alias,
|
||||
"max_budget": budget,
|
||||
"budget_duration": "30d",
|
||||
"models": ["deepseek-v4-flash", "gpt-oss-120b"],
|
||||
|
||||
Reference in new issue
Block a user