forked from phoenix/litellm-mirror
fix(proxy_server.py): fix model info check
This commit is contained in:
parent
5bcf92f4f5
commit
c81c9c2583
2 changed files with 46 additions and 1 deletions
|
@ -6409,6 +6409,9 @@ async def add_new_model(model_params: ModelParams):
|
||||||
async def model_info_v2(
|
async def model_info_v2(
|
||||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
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
|
global llm_model_list, general_settings, user_config_file_path, proxy_config
|
||||||
|
|
||||||
# Load existing config
|
# Load existing config
|
||||||
|
@ -6550,7 +6553,7 @@ async def model_info_v1(
|
||||||
|
|
||||||
if len(user_api_key_dict.models) > 0:
|
if len(user_api_key_dict.models) > 0:
|
||||||
model_names = user_api_key_dict.models
|
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:
|
else:
|
||||||
all_models = config["model_list"]
|
all_models = config["model_list"]
|
||||||
for model in all_models:
|
for model in all_models:
|
||||||
|
|
|
@ -280,6 +280,29 @@ async def get_key_info(session, call_key, get_key=None):
|
||||||
return await response.json()
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_key_info():
|
async def test_key_info():
|
||||||
"""
|
"""
|
||||||
|
@ -305,6 +328,25 @@ async def test_key_info():
|
||||||
assert status == 403
|
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):
|
async def get_spend_logs(session, request_id):
|
||||||
url = f"http://0.0.0.0:4000/spend/logs?request_id={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"}
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue