fix raise correct error 404 when /key/info is called on non-existent key (#6653)

* fix raise correct error on /key/info

* add not_found_error error

* fix key not found in DB error

* use 1 helper for checking token hash

* fix error code on key info

* fix test key gen prisma

* test_generate_and_call_key_info

* test fix test_call_with_valid_model_using_all_models

* fix key info tests
This commit is contained in:
Ishaan Jaff 2024-11-11 21:00:39 -08:00 committed by GitHub
parent 25bae4cc23
commit de2f9aed3a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 3593 additions and 57 deletions

View file

@ -1424,9 +1424,7 @@ class PrismaClient:
# check if plain text or hash
if token is not None:
if isinstance(token, str):
hashed_token = token
if token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
hashed_token = _hash_token_if_needed(token=token)
verbose_proxy_logger.debug(
f"PrismaClient: find_unique for token: {hashed_token}"
)
@ -1493,8 +1491,7 @@ class PrismaClient:
if token is not None:
where_filter["token"] = {}
if isinstance(token, str):
if token.startswith("sk-"):
token = self.hash_token(token=token)
token = _hash_token_if_needed(token=token)
where_filter["token"]["in"] = [token]
elif isinstance(token, list):
hashed_tokens = []
@ -1630,9 +1627,7 @@ class PrismaClient:
# check if plain text or hash
if token is not None:
if isinstance(token, str):
hashed_token = token
if token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
hashed_token = _hash_token_if_needed(token=token)
verbose_proxy_logger.debug(
f"PrismaClient: find_unique for token: {hashed_token}"
)
@ -1912,8 +1907,7 @@ class PrismaClient:
if token is not None:
print_verbose(f"token: {token}")
# check if plain text or hash
if token.startswith("sk-"):
token = self.hash_token(token=token)
token = _hash_token_if_needed(token=token)
db_data["token"] = token
response = await self.db.litellm_verificationtoken.update(
where={"token": token}, # type: ignore
@ -2424,6 +2418,18 @@ def hash_token(token: str):
return hashed_token
def _hash_token_if_needed(token: str) -> str:
"""
Hash the token if it's a string and starts with "sk-"
Else return the token as is
"""
if token.startswith("sk-"):
return hash_token(token=token)
else:
return token
def _extract_from_regex(duration: str) -> Tuple[int, str]:
match = re.match(r"(\d+)(mo|[smhd]?)", duration)