This commit is contained in:
Krrish Dholakia 2024-08-08 10:30:15 -07:00
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

View file

@ -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"]}]

View file

@ -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."
)

View file

@ -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:

View file

@ -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)}")