diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b15d6ba1..413d7aa4f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6409,6 +6409,9 @@ async def add_new_model(model_params: ModelParams): async def model_info_v2( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): + """ + BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. + """ global llm_model_list, general_settings, user_config_file_path, proxy_config # Load existing config @@ -6550,7 +6553,7 @@ async def model_info_v1( if len(user_api_key_dict.models) > 0: model_names = user_api_key_dict.models - all_models = [m for m in config["model_list"] if m in model_names] + all_models = [m for m in config["model_list"] if m["model_name"] in model_names] else: all_models = config["model_list"] for model in all_models: diff --git a/tests/test_keys.py b/tests/test_keys.py index eb1c55ca3..4a563cc18 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -280,6 +280,29 @@ async def get_key_info(session, call_key, get_key=None): return await response.json() +async def get_model_info(session, call_key): + """ + Make sure only models user has access to are returned + """ + url = "http://0.0.0.0:4000/model/info" + headers = { + "Authorization": f"Bearer {call_key}", + "Content-Type": "application/json", + } + + async with session.get(url, headers=headers) as response: + status = response.status + response_text = await response.text() + print(response_text) + print() + + if status != 200: + raise Exception( + f"Request did not return a 200 status code: {status}. Responses {response_text}" + ) + return await response.json() + + @pytest.mark.asyncio async def test_key_info(): """ @@ -305,6 +328,25 @@ async def test_key_info(): assert status == 403 +@pytest.mark.asyncio +async def test_model_info(): + """ + Get model info for models key has access to + """ + async with aiohttp.ClientSession() as session: + key_gen = await generate_key(session=session, i=0) + key = key_gen["key"] + # as admin # + admin_models = await get_model_info(session=session, call_key="sk-1234") + admin_models = admin_models["data"] + # as key itself # + user_models = await get_model_info(session=session, call_key=key) + user_models = user_models["data"] + + assert len(admin_models) > len(user_models) + assert len(user_models) > 0 + + async def get_spend_logs(session, request_id): url = f"http://0.0.0.0:4000/spend/logs?request_id={request_id}" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}