phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(fix) dynamoDB + hashed tokens bug

Browse source
This commit is contained in:
ishaan-jaff 2024-01-29 11:17:27 -08:00
parent 3836ba4e65
commit e672e5b928
2 changed files with 6 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/db/dynamo_db.py
View file

@ -187,6 +187,7 @@ class DynamoDBWrapper(CustomDB):
elif table_name == "spend": elif table_name == "spend":
table = client.table(self.database_arguments.spend_table_name) table = client.table(self.database_arguments.spend_table_name)
value = value.copy()
for k, v in value.items(): for k, v in value.items():
if k == "token" and value[k].startswith("sk-"): if k == "token" and value[k].startswith("sk-"):
value[k] = hash_token(token=v) value[k] = hash_token(token=v)
@ -311,6 +312,7 @@ class DynamoDBWrapper(CustomDB):
# Initialize an empty UpdateExpression # Initialize an empty UpdateExpression
actions: List = [] actions: List = []
value = value.copy()
for k, v in value.items(): for k, v in value.items():
# Convert datetime object to ISO8601 string # Convert datetime object to ISO8601 string
if isinstance(v, datetime): if isinstance(v, datetime):

4
litellm/tests/test_key_generate_dynamodb.py
View file

@ -116,6 +116,10 @@ def test_call_with_invalid_key(custom_db_client):
def test_call_with_invalid_model(custom_db_client): def test_call_with_invalid_model(custom_db_client):
# 3. Make a call to a key with an invalid model - expect to fail # 3. Make a call to a key with an invalid model - expect to fail
from litellm._logging import verbose_proxy_logger
import logging
verbose_proxy_logger.setLevel(logging.DEBUG)
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client) setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
try: try: