phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat(proxy_server.py): support key lists for /key/info

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-10 14:24:12 -08:00
parent 525b0dc4fd
commit 6b83001459
4 changed files with 61 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

40
litellm/proxy/proxy_server.py
View file

@ -2894,7 +2894,7 @@ async def update_key_fn(request: Request, data: UpdateKeyRequest):
@router.post(
"/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
)
async def delete_key_fn(data: DeleteKeyRequest):
async def delete_key_fn(data: KeyRequest):
"""
Delete a key from the key management system.
@ -2958,6 +2958,7 @@ async def delete_key_fn(data: DeleteKeyRequest):
"/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
)
async def info_key_fn(
data: Optional[KeyRequest] = None,
key: Optional[str] = fastapi.Query(
default=None, description="Key in the request parameters"
),
@ -2966,16 +2967,23 @@ async def info_key_fn(
"""
Retrieve information about a key.
Parameters:
key: Optional[str] = Query parameter representing the key in the request
key: Optional[str | list] = Query parameter representing the key(s) in the request
user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key
Returns:
Dict containing the key and its associated information
Example Curl:
**1 key**
```
curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \
-H "Authorization: Bearer sk-1234"
```
**multiple keys**
```
curl -X GET "http://0.0.0.0:8000/key/info" \
-H "Authorization: Bearer sk-1234" \
-d {"keys": ["sk-1", "sk-2", "sk-3"]}
```
Example Curl - if no key is passed, it will use the Key Passed in Authorization Header
```
@ -2989,16 +2997,24 @@ async def info_key_fn(
raise Exception(
f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
if key == None:
key = user_api_key_dict.api_key
key_info = await prisma_client.get_data(token=key)
## REMOVE HASHED TOKEN INFO BEFORE RETURNING ##
try:
key_info = key_info.model_dump() # noqa
except:
# if using pydantic v1
key_info = key_info.dict()
key_info.pop("token")
if key is None and data is None:
key_info = await prisma_client.get_data(token=user_api_key_dict.api_key)
if key is not None: # single key
key_info = await prisma_client.get_data(token=key)
elif data is not None: # list of keys
key_info = await prisma_client.get_data(
token=data.keys, table_name="key", query_type="find_all"
)
filtered_key_info = []
for k in key_info:
try:
k = k.model_dump() # noqa
except:
# if using pydantic v1
k = k.dict()
filtered_key_info.append(k)
return {"key": data.keys, "info": filtered_key_info}
return {"key": key, "info": key_info}
except Exception as e:
if isinstance(e, HTTPException):