(fix) dynamo - use hashed_api keys in table

This commit is contained in:
ishaan-jaff 2024-01-26 14:32:00 -08:00
parent 0ee8e8f081
commit 2c79585e18

View file

@ -5,6 +5,7 @@ from litellm.proxy._types import (
LiteLLM_Config,
LiteLLM_UserTable,
)
from litellm.proxy.utils import hash_token
from litellm import get_secret
from typing import Any, List, Literal, Optional, Union
import json
@ -187,6 +188,8 @@ class DynamoDBWrapper(CustomDB):
table = client.table(self.database_arguments.spend_table_name)
for k, v in value.items():
if k == "token" and value[k].startswith("sk-"):
value[k] = hash_token(token=v)
if isinstance(v, datetime):
value[k] = v.isoformat()
@ -229,6 +232,10 @@ class DynamoDBWrapper(CustomDB):
table = client.table(self.database_arguments.config_table_name)
key_name = "param_name"
if key_name == "token" and key.startswith("sk-"):
# ensure it's hashed
key = hash_token(token=key)
response = await table.get_item({key_name: key})
new_response: Any = None
@ -308,6 +315,8 @@ class DynamoDBWrapper(CustomDB):
# Convert datetime object to ISO8601 string
if isinstance(v, datetime):
v = v.isoformat()
if k == "token" and value[k].startswith("sk-"):
value[k] = hash_token(token=v)
# Accumulate updates
actions.append((F(k), Value(value=v)))