diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 3b73f65c917..415d5d546e9 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -557,6 +557,7 @@ def prepare_call( litellm_params={ "litellm_call_id": kwargs.get("litellm_call_id"), "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model), + **({} if (model_info := kwargs.get("model_info")) is None else {"model_info": model_info}), }, custom_llm_provider=provider, ) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index e5501310556..a8c2d4e7805 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -6,6 +6,8 @@ import copy import json import re import threading +import uuid +from datetime import datetime from collections.abc import Mapping, Sequence from concurrent.futures import wait from types import MappingProxyType @@ -296,6 +298,43 @@ async def test_router_dispatches_typesafe_decisions_without_api_base( assert response.answers["sentiment"].choice == "positive" +@pytest.mark.asyncio +async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_finishes( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "jev", + "litellm_params": {"model": "typesafe/jev-latest", "api_key": "k"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("jev-a", "jev-b") + ], + routing_strategy="least-busy", + ) + + async def call_through_proxy_logging() -> object: + logging_obj, _ = litellm.utils.function_setup( + original_function="adecisions", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="jev", + litellm_call_id=str(uuid.uuid4()), + ) + response: Final = await router.adecisions( + model="jev", + input="Is this a defect?", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + litellm_logging_obj=logging_obj, + ) + await _drain_logging_worker() + return response._hidden_params["model_id"] + + assert (await call_through_proxy_logging(), await call_through_proxy_logging()) == ("jev-a", "jev-a") + + def test_decisions_uses_the_same_wire_contract_for_sync_calls(respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 46964bdb9f8..3dbd90b814b 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -236,9 +236,7 @@ class OpenBreakerRedis(SharedRedisCounters): @pytest.mark.asyncio -async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( - caplog: pytest.LogCaptureFixture, -) -> None: +async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None: worker: Final = _worker(OpenBreakerRedis()) with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): @@ -410,7 +408,9 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts if use_async else router.get_available_deployment(model=ROUTER_GROUP, messages=None) ) - response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) + response: Final = await router.acompletion( + model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}] + ) for _ in range(10): await asyncio.sleep(0) await GLOBAL_LOGGING_WORKER.flush()