forked from phoenix/litellm-mirror
fix(user_api_key_auth.py): Fixes https://github.com/BerriAI/litellm/issues/5111
This commit is contained in:
parent
5d664c0441
commit
5703da9b42
7 changed files with 37 additions and 23 deletions
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
@ -1,15 +1,4 @@
|
|||
model_list:
|
||||
- model_name: "gpt-3.5-turbo"
|
||||
litellm_params:
|
||||
model: "gpt-3.5-turbo"
|
||||
- model_name: "gpt-4"
|
||||
litellm_params:
|
||||
model: "gpt-4"
|
||||
api_key: "bad_key"
|
||||
- model_name: "gpt-4o"
|
||||
litellm_params:
|
||||
model: "gpt-4o"
|
||||
|
||||
litellm_settings:
|
||||
enable_json_schema_validation: true
|
||||
fallbacks: [{"gpt-3.5-turbo": ["gpt-4", "gpt-4o"]}]
|
||||
|
|
|
@ -375,7 +375,7 @@ async def get_user_object(
|
|||
await user_api_key_cache.async_set_cache(key=user_id, value=_response)
|
||||
|
||||
return _response
|
||||
except Exception as e: # if user not in db
|
||||
except Exception: # if user not in db
|
||||
raise ValueError(
|
||||
f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call."
|
||||
)
|
||||
|
|
|
@ -652,14 +652,22 @@ async def user_api_key_auth(
|
|||
|
||||
# Check 2. If user_id for this token is in budget - done in common_checks()
|
||||
if valid_token.user_id is not None:
|
||||
user_obj = await get_user_object(
|
||||
user_id=valid_token.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
try:
|
||||
user_obj = await get_user_object(
|
||||
user_id=valid_token.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - {}".format(
|
||||
str(e)
|
||||
)
|
||||
)
|
||||
user_obj = None
|
||||
|
||||
# Check 3. Check if user is in their team budget
|
||||
if valid_token.team_member_spend is not None:
|
||||
|
|
|
@ -798,3 +798,23 @@ async def test_key_model_list(model_access, model_access_level, model_endpoint):
|
|||
elif model_endpoint == "/model/info":
|
||||
assert isinstance(model_list["data"], list)
|
||||
assert len(model_list["data"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_user_not_in_db():
|
||||
"""
|
||||
- Create a key with unique user-id (not in db)
|
||||
- Check if key can make `/chat/completion` call
|
||||
"""
|
||||
my_unique_user = str(uuid.uuid4())
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(
|
||||
session=session,
|
||||
i=0,
|
||||
user_id=my_unique_user,
|
||||
)
|
||||
key = key_gen["key"]
|
||||
try:
|
||||
await chat_completion(session=session, key=key)
|
||||
except Exception as e:
|
||||
pytest.fail(f"Expected this call to work - {str(e)}")
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue