phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(router.py): support request prioritization for text completion c… (#7540)

Browse source 
* feat(router.py): support request prioritization for text completion calls

* fix(internal_user_endpoints.py): fix sql query to return all keys, including null team id keys on `/user/info`

Fixes https://github.com/BerriAI/litellm/issues/7485

* fix: fix linting errors

* fix: fix linting error

* test(test_router_helper_utils.py): add direct test for '_schedule_factory'

Fixes code qa test
This commit is contained in:
Krish Dholakia 2025-01-03 19:35:44 -08:00 • committed by GitHub
parent fb1272b46b
commit db82b3bb2a
7 changed files with 229 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

43
tests/router_unit_tests/test_router_helper_utils.py
View file

@ -166,6 +166,49 @@ async def test_router_schedule_acompletion(model_list):
assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!"
@pytest.mark.asyncio
async def test_router_schedule_atext_completion(model_list):
"""Test if the 'schedule_atext_completion' function is working correctly"""
from litellm.types.utils import TextCompletionResponse
router = Router(model_list=model_list)
with patch.object(
router, "_atext_completion", AsyncMock()
) as mock_atext_completion:
mock_atext_completion.return_value = TextCompletionResponse()
response = await router.atext_completion(
model="gpt-3.5-turbo",
prompt="Hello, how are you?",
priority=1,
)
mock_atext_completion.assert_awaited_once()
assert "priority" not in mock_atext_completion.call_args.kwargs
@pytest.mark.asyncio
async def test_router_schedule_factory(model_list):
"""Test if the 'schedule_atext_completion' function is working correctly"""
from litellm.types.utils import TextCompletionResponse
router = Router(model_list=model_list)
with patch.object(
router, "_atext_completion", AsyncMock()
) as mock_atext_completion:
mock_atext_completion.return_value = TextCompletionResponse()
response = await router._schedule_factory(
model="gpt-3.5-turbo",
args=(
"gpt-3.5-turbo",
"Hello, how are you?",
),
priority=1,
kwargs={},
original_function=router.atext_completion,
)
mock_atext_completion.assert_awaited_once()
assert "priority" not in mock_atext_completion.call_args.kwargs
@pytest.mark.asyncio
async def test_router_arealtime(model_list):
"""Test if the '_arealtime' function is working correctly"""