add test_chat_completion_with_no_tags

This commit is contained in:
Ishaan Jaff 2024-08-29 13:54:11 -07:00
parent 84bda9cc80
commit f592aeaa38
3 changed files with 34 additions and 6 deletions

View file

@ -4,6 +4,17 @@ model_list:
model: openai/fake model: openai/fake
api_key: fake-key api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/ api_base: https://exampleopenaiendpoint-production.up.railway.app/
tags: ["teamA"]
model_info:
id: "teama"
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
tags: ["teamB"]
model_info:
id: "teamb"
- model_name: rerank-english-v3.0 - model_name: rerank-english-v3.0
litellm_params: litellm_params:
model: cohere/rerank-english-v3.0 model: cohere/rerank-english-v3.0

View file

@ -1,15 +1,21 @@
model_list: model_list:
- model_name: gpt-4 - model_name: fake-openai-endpoint
litellm_params: litellm_params:
model: openai/fake model: openai/fake
api_key: fake-key api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/ api_base: https://exampleopenaiendpoint-production.up.railway.app/
tags: ["teamA"] # 👈 Key Change tags: ["teamA"]
- model_name: gpt-4 model_info:
id: "teama"
- model_name: fake-openai-endpoint
litellm_params: litellm_params:
model: openai/gpt-4o model: openai/fake
api_key: os.environ/OPENAI_API_KEY api_key: fake-key
tags: ["teamB"] # 👈 Key Change api_base: https://exampleopenaiendpoint-production.up.railway.app/
tags: ["teamB"]
model_info:
id: "teamb"
router_settings: router_settings:
enable_tag_filtering: True # 👈 Key Change enable_tag_filtering: True # 👈 Key Change

View file

@ -119,3 +119,14 @@ async def test_team_tag_routing():
print(response_b) print(response_b)
print(headers) print(headers)
assert headers["x-litellm-model-id"] == "teamb", "Model ID should be teamB" assert headers["x-litellm-model-id"] == "teamb", "Model ID should be teamB"
@pytest.mark.asyncio()
async def test_chat_completion_with_no_tags():
async with aiohttp.ClientSession() as session:
key = LITELLM_MASTER_KEY
response, headers = await chat_completion(session, key)
headers = dict(headers)
print(response)
print(headers)
assert response is not None