mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(health): log Test Connection spend with the mode's call type and model id (#50987)
* fix(health): price connection tests using their resolved mode Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(health): avoid deployment pricing on alternate probes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(health): preserve pricing for equivalent model names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): avoid authentication when comparing health check models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover health probe pricing without a deployment model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): pin the clock in the rpm rate limit test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kenneth-berri
GitHub
parent
59538ff1d9
commit
b4227652c3
7 changed files with 392 additions and 29 deletions
|
|
@ -9,6 +9,8 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.base_llm.ocr.transformation import DocumentType
|
||||
from litellm.rust_bridge import runtime
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR_HEALTH_CHECK_DOCUMENT
|
||||
|
|
@ -190,7 +192,10 @@ class HealthCheckHelpers:
|
|||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
|
||||
_metadata_variable_name: Final = "litellm_metadata"
|
||||
litellm_metadata: Final = HealthCheckHelpers._get_metadata_for_health_check_call()
|
||||
litellm_metadata: Final[Mapping[str, object]] = {
|
||||
**HealthCheckHelpers._get_metadata_for_health_check_call(),
|
||||
"model_info": TypeAdapter(dict[str, object]).validate_python(model_params).get("model_info") or {},
|
||||
}
|
||||
model_params[_metadata_variable_name] = litellm_metadata
|
||||
model_params = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
|
||||
data=model_params,
|
||||
|
|
|
|||
|
|
@ -9512,6 +9512,7 @@ async def ahealth_check(
|
|||
from litellm.litellm_core_utils.cached_imports import get_litellm_logging_class
|
||||
from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers, default_health_check_mode
|
||||
from litellm.litellm_core_utils.health_check_utils import OPTIONAL_STR
|
||||
from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
|
||||
|
||||
# Use cached import helper to lazy-load Logging class (only loads when function is called)
|
||||
Logging: Final = get_litellm_logging_class()
|
||||
|
|
@ -9556,6 +9557,24 @@ async def ahealth_check(
|
|||
mode = mode or default_health_check_mode(
|
||||
requested_model=requested_model, model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
litellm_logging_obj.model = requested_model
|
||||
litellm_logging_obj.call_type = {
|
||||
"completion": "atext_completion",
|
||||
"embedding": "aembedding",
|
||||
"audio_speech": "aspeech",
|
||||
"audio_transcription": "atranscription",
|
||||
"image_generation": "aimage_generation",
|
||||
"image_edit": "aimage_edit",
|
||||
"video_generation": "avideo_generation",
|
||||
"rerank": "arerank",
|
||||
"realtime": "arealtime",
|
||||
"responses": "aresponses",
|
||||
"anthropic_messages": "anthropic_messages",
|
||||
"ocr": "aocr",
|
||||
"decisions": "asystemone",
|
||||
"extraction": "aextraction",
|
||||
"batch": ("alist_batches" if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS else "acompletion"),
|
||||
}.get("decisions" if is_decisions_model_mode(mode) else mode, "acompletion")
|
||||
if "*" in model:
|
||||
return await HealthCheckHelpers.ahealth_check_wildcard_models(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm.constants import (
|
|||
DEFAULT_HEALTH_CHECK_PROMPT,
|
||||
HEALTH_CHECK_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
|
||||
from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode
|
||||
from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin, is_model_capability_field
|
||||
from litellm.router_strategy.auto_router.auto_router_model_naming import (
|
||||
|
|
@ -831,6 +832,24 @@ _PROBE_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object])
|
|||
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _same_health_check_model(deployment_model: str | None, probe_params: Mapping[str, object]) -> bool:
|
||||
model: Final = probe_params.get("model")
|
||||
if model == deployment_model:
|
||||
return True
|
||||
if not isinstance(model, str) or deployment_model is None:
|
||||
return False
|
||||
provider_param: Final = probe_params.get("custom_llm_provider")
|
||||
provider: Final = provider_param if isinstance(provider_param, str) else None
|
||||
try:
|
||||
deployment_name, deployment_provider = resolve_model_and_provider_for_metadata(
|
||||
model=deployment_model, custom_llm_provider=provider
|
||||
)
|
||||
probe_name, probe_provider = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=provider)
|
||||
except litellm.exceptions.BadRequestError:
|
||||
return False
|
||||
return (deployment_name, deployment_provider) == (probe_name, probe_provider)
|
||||
|
||||
|
||||
def update_litellm_params_for_health_check(
|
||||
model_info: dict, litellm_params: dict, deployment_model: str | None = None
|
||||
) -> dict:
|
||||
|
|
@ -895,6 +914,15 @@ def update_litellm_params_for_health_check(
|
|||
if mode == "audio_speech":
|
||||
litellm_params["voice"] = model_info.get("health_check_voice", "alloy")
|
||||
|
||||
logging_model_info: Final = _MODEL_INFO_ADAPTER.validate_python(
|
||||
_PROBE_PARAMS_ADAPTER.validate_python(litellm_params).get("model_info") or {}
|
||||
)
|
||||
probe_logging_model_info: Final = (
|
||||
logging_model_info
|
||||
if _same_health_check_model(base_model_owner, _PROBE_PARAMS_ADAPTER.validate_python(litellm_params))
|
||||
else {key: value for key, value in logging_model_info.items() if key == "id"}
|
||||
)
|
||||
|
||||
# Handle Bedrock region routing format: bedrock/region/model
|
||||
# This is needed because health checks bypass get_llm_provider() for the model param
|
||||
# Issue #15807: Without this, health checks send "region/model" as the model ID to AWS
|
||||
|
|
@ -941,7 +969,12 @@ def update_litellm_params_for_health_check(
|
|||
model_id_params: Final = {} if model_id_pin is None else {"model_id": model_id_pin}
|
||||
return {
|
||||
key: value
|
||||
for key, value in {**audio_output_params, **model_id_params, **probe_params}.items()
|
||||
for key, value in {
|
||||
**audio_output_params,
|
||||
**model_id_params,
|
||||
**probe_params,
|
||||
**({"model_info": probe_logging_model_info} if "model_info" in probe_params else {}),
|
||||
}.items()
|
||||
if not is_model_capability_field(key)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -2695,7 +2695,10 @@ async def test_model_connection(
|
|||
stored_model: Final = _OBJECT_MAPPING.validate_python(config_litellm_params).get("model")
|
||||
litellm_params = update_litellm_params_for_health_check(
|
||||
model_info=dict(probe_model_info),
|
||||
litellm_params=litellm_params,
|
||||
litellm_params={
|
||||
**litellm_params,
|
||||
"model_info": _OBJECT_MAPPING.validate_python(resolved_model_info or {}),
|
||||
},
|
||||
deployment_model=stored_model if isinstance(stored_model, str) else None,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,10 +4,19 @@ from datetime import datetime, timedelta, timezone
|
|||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from pydantic import TypeAdapter
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.chatgpt import authenticator as chatgpt_authenticator
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
|
||||
from litellm.proxy.health_check import _same_health_check_model, update_litellm_params_for_health_check
|
||||
from litellm.proxy.health_endpoints._health_endpoints import (
|
||||
_aggregate_health_check_results,
|
||||
_build_model_param_to_info_mapping,
|
||||
|
|
@ -16,8 +25,10 @@ from litellm.proxy.health_endpoints._health_endpoints import (
|
|||
_save_health_check_results_if_changed,
|
||||
_save_health_check_to_db,
|
||||
latest_health_checks_endpoint,
|
||||
test_model_connection as probe_model_connection,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -946,5 +957,213 @@ def test_filter_deployments_to_model_groups():
|
|||
assert filter_deployments_to_model_groups(model_list, frozenset()) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"pricing", [{"base_model": "gpt-5.4"}, {"input_cost_per_token": 0.02, "output_cost_per_token": 0.03}]
|
||||
)
|
||||
@pytest.mark.parametrize("deployment_id", [None, "alternate-probe-deployment"])
|
||||
@pytest.mark.parametrize("override", ["health_check_model", "health_check_params", "same_model"])
|
||||
async def test_alternate_health_check_model_uses_its_own_price(
|
||||
pricing: dict[str, object], deployment_id: str | None, override: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
events: Final[asyncio.Queue[StandardLoggingPayload]] = asyncio.Queue()
|
||||
|
||||
class SpendCapture(CustomLogger):
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
events.put_nowait(TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]))
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [SpendCapture()])
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"prisma_client",
|
||||
PrismaClient(
|
||||
database_url="postgresql://test:test@localhost/test", proxy_logging_obj=proxy_server.proxy_logging_obj
|
||||
),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"health-probe",
|
||||
{
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
},
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"gpt-5.4",
|
||||
{"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "litellm_provider": "openai", "mode": "chat"},
|
||||
)
|
||||
if deployment_id is not None:
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost, deployment_id, {**pricing, "litellm_provider": "openai", "mode": "chat"}
|
||||
)
|
||||
params: Final = {"api_base": "https://health-test.example/v1", "api_key": "test-key"}
|
||||
model_info: Final = {
|
||||
**pricing,
|
||||
**(
|
||||
{"health_check_params": {"model": "openai/health-probe"}}
|
||||
if override == "health_check_params"
|
||||
else {"health_check_model": "openai/health-probe"}
|
||||
),
|
||||
**({"id": deployment_id} if deployment_id is not None else {}),
|
||||
}
|
||||
payload: Final = {
|
||||
"id": "chatcmpl-health-probe",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "health-probe",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
with respx.mock:
|
||||
respx.post("https://health-test.example/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json=payload)
|
||||
)
|
||||
await litellm.acompletion(
|
||||
model="openai/health-probe",
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
litellm_metadata={"model_info": model_info} if override == "same_model" else {},
|
||||
**params,
|
||||
)
|
||||
normal: Final = await asyncio.wait_for(events.get(), timeout=10)
|
||||
result: Final = await probe_model_connection(
|
||||
request=Request({"type": "http", "headers": []}),
|
||||
mode="chat",
|
||||
litellm_params={"model": "openai/health-probe" if override == "same_model" else "openai/gpt-5.4", **params},
|
||||
model_info=model_info,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
assert result["status"] == "success", result
|
||||
health: Final = await asyncio.wait_for(events.get(), timeout=10)
|
||||
|
||||
assert health["call_type"] == normal["call_type"]
|
||||
assert health["model"] == normal["model"]
|
||||
assert health["response_cost"] == normal["response_cost"]
|
||||
assert health["response_cost"] > 0
|
||||
assert health["model_id"] == (deployment_id or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("stored_model", "requested_model"),
|
||||
[("openai/health-probe", "health-probe"), ("health-probe", "openai/health-probe")],
|
||||
)
|
||||
async def test_equivalent_health_check_model_names_keep_deployment_prices(
|
||||
stored_model: str, requested_model: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
events: Final[asyncio.Queue[StandardLoggingPayload]] = asyncio.Queue()
|
||||
|
||||
class SpendCapture(CustomLogger):
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
events.put_nowait(TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]))
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [SpendCapture()])
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"health-probe",
|
||||
{
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
},
|
||||
)
|
||||
params: Final = {
|
||||
"api_base": "https://equivalent-health.example/v1",
|
||||
"api_key": "test-key",
|
||||
"custom_llm_provider": "openai",
|
||||
}
|
||||
deployment_id: Final = "priced-health-deployment"
|
||||
model_info: Final = {
|
||||
"id": deployment_id,
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 0.02,
|
||||
"output_cost_per_token": 0.03,
|
||||
}
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "priced-health",
|
||||
"litellm_params": {"model": stored_model, **params},
|
||||
"model_info": model_info,
|
||||
}
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"prisma_client",
|
||||
PrismaClient(
|
||||
database_url="postgresql://test:test@localhost/test", proxy_logging_obj=proxy_server.proxy_logging_obj
|
||||
),
|
||||
)
|
||||
payload: Final = {
|
||||
"id": "chatcmpl-equivalent-health",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "health-probe",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
with respx.mock:
|
||||
respx.post("https://equivalent-health.example/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json=payload)
|
||||
)
|
||||
result: Final = await probe_model_connection(
|
||||
request=Request({"type": "http", "headers": []}),
|
||||
mode="chat",
|
||||
litellm_params={"model": requested_model, "custom_llm_provider": "openai"},
|
||||
model_info={"id": deployment_id},
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
assert result["status"] == "success", result
|
||||
health: Final = await asyncio.wait_for(events.get(), timeout=10)
|
||||
|
||||
assert health["response_cost"] == 10 * model_info["input_cost_per_token"] + 5 * model_info["output_cost_per_token"]
|
||||
assert health["model"] == "health-probe"
|
||||
assert health["model_id"] == deployment_id
|
||||
|
||||
|
||||
def test_health_check_without_deployment_model_keeps_probe_but_drops_deployment_prices() -> None:
|
||||
deployment_info: Final = {
|
||||
"id": "probe-deployment",
|
||||
"base_model": "openai/deployment-model",
|
||||
"input_cost_per_token": 0.02,
|
||||
"output_cost_per_token": 0.03,
|
||||
}
|
||||
probe_model: Final = "openai/health-probe"
|
||||
|
||||
probe_params: Final = update_litellm_params_for_health_check(
|
||||
model_info={"mode": "chat", "health_check_model": probe_model},
|
||||
litellm_params={"model_info": deployment_info},
|
||||
)
|
||||
|
||||
assert probe_params["model"] == probe_model
|
||||
assert probe_params["model_info"] == {"id": deployment_info["id"]}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("model", "expected"), [("health-probe", True), ("other-health-probe", False)])
|
||||
def test_chatgpt_health_check_model_comparison_does_not_access_credentials(
|
||||
model: str, expected: bool, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
def forbidden_auth_file_access(*args: object, **kwargs: object) -> None:
|
||||
pytest.fail("Health-check model comparison accessed the ChatGPT auth file")
|
||||
|
||||
monkeypatch.setattr(chatgpt_authenticator, "open", forbidden_auth_file_access, raising=False)
|
||||
|
||||
assert (
|
||||
_same_health_check_model("chatgpt/health-probe", {"model": model, "custom_llm_provider": "chatgpt"}) is expected
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ from litellm.types.utils import (
|
|||
Delta,
|
||||
ModelResponseStream,
|
||||
StandardCallbackDynamicParams,
|
||||
StandardLoggingPayload,
|
||||
StreamingChoices,
|
||||
TextCompletionResponse,
|
||||
Usage,
|
||||
|
|
@ -103,6 +104,63 @@ def preserve_litellm_completion_state(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["image_generation", "embedding"])
|
||||
@pytest.mark.parametrize("deployment_id", [None, "health-deployment"])
|
||||
async def test_health_check_logs_the_same_cost_and_deployment_as_normal_calls(
|
||||
mode: str, deployment_id: str | None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
events: Final[asyncio.Queue[StandardLoggingPayload]] = asyncio.Queue()
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"standard/1024-x-1024/health-image",
|
||||
{"input_cost_per_image": 0.25, "litellm_provider": "openai", "mode": "image_generation"},
|
||||
)
|
||||
|
||||
class SpendCapture(CustomLogger):
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
events.put_nowait(TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]))
|
||||
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [SpendCapture()])
|
||||
model: Final = "openai/health-image" if mode == "image_generation" else "openai/text-embedding-3-small"
|
||||
payload: Final = (
|
||||
{"created": 1, "data": [{"url": "https://example.com/image.png"}]}
|
||||
if mode == "image_generation"
|
||||
else {
|
||||
"object": "list",
|
||||
"data": [{"object": "embedding", "embedding": [0.1], "index": 0}],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": {"prompt_tokens": 10, "total_tokens": 10},
|
||||
}
|
||||
)
|
||||
params: Final = {
|
||||
"model": model,
|
||||
"api_base": "https://health-test.example/v1",
|
||||
"api_key": "test-key",
|
||||
**({"model_info": {"id": deployment_id}} if deployment_id else {}),
|
||||
**({"size": "1024x1024", "quality": "standard", "n": 1} if mode == "image_generation" else {}),
|
||||
}
|
||||
with respx.mock:
|
||||
respx.post(url__regex=r"https://health-test.example/v1/.*").mock(return_value=httpx.Response(200, json=payload))
|
||||
if mode == "image_generation":
|
||||
await litellm.aimage_generation(**params, prompt="test")
|
||||
else:
|
||||
await litellm.aembedding(**params, input=["test"])
|
||||
normal: Final = await asyncio.wait_for(events.get(), timeout=10)
|
||||
result: Final = await litellm.ahealth_check(model_params=params.copy(), mode=mode, prompt="test")
|
||||
assert "error" not in result, result
|
||||
health: Final = await asyncio.wait_for(events.get(), timeout=10)
|
||||
|
||||
assert health["call_type"] == normal["call_type"]
|
||||
assert health["model"] == normal["model"]
|
||||
assert health["response_cost"] == normal["response_cost"]
|
||||
assert health["response_cost"] > 0
|
||||
assert health["model_id"] == (deployment_id or "")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
|
|
@ -1,11 +1,43 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from litellm.utils import calculate_max_parallel_requests
|
||||
|
||||
_ROUTER_CLOCK: Final = "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime"
|
||||
_FIRST_MINUTE: Final = datetime.datetime(2026, 1, 1, 12, 0, 59, tzinfo=datetime.timezone.utc)
|
||||
_NEXT_MINUTE: Final = datetime.datetime(2026, 1, 1, 12, 1, 0, tzinfo=datetime.timezone.utc)
|
||||
|
||||
|
||||
def _rpm_limited_router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "limited",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "test-key",
|
||||
"rpm": 1,
|
||||
},
|
||||
}
|
||||
],
|
||||
routing_strategy="usage-based-routing-v2",
|
||||
enable_pre_call_checks=True,
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
_RPM_LIMITED_REQUEST: Final = {
|
||||
"model": "limited",
|
||||
"messages": [{"role": "user", "content": "rate limit"}],
|
||||
"mock_response": "response",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("max_parallel_requests", "tpm", "rpm", "default_max_parallel_requests", "expected"),
|
||||
|
|
@ -37,32 +69,26 @@ def test_scenario(
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_parallel_requests_rpm_rate_limiting() -> None:
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "limited",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "test-key",
|
||||
"rpm": 1,
|
||||
},
|
||||
}
|
||||
],
|
||||
routing_strategy="usage-based-routing-v2",
|
||||
enable_pre_call_checks=True,
|
||||
num_retries=0,
|
||||
)
|
||||
request: Final = {
|
||||
"model": "limited",
|
||||
"messages": [{"role": "user", "content": "rate limit"}],
|
||||
"mock_response": "response",
|
||||
}
|
||||
router: Final = _rpm_limited_router()
|
||||
|
||||
first_response: Final = await router.acompletion(**request)
|
||||
with patch(_ROUTER_CLOCK, return_value=_FIRST_MINUTE):
|
||||
first_response: Final = await router.acompletion(**_RPM_LIMITED_REQUEST)
|
||||
|
||||
assert first_response.choices[0].message.content == "response"
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(**request)
|
||||
assert first_response.choices[0].message.content == "response"
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(**_RPM_LIMITED_REQUEST)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rpm_limit_resets_in_the_next_minute() -> None:
|
||||
router: Final = _rpm_limited_router()
|
||||
|
||||
with patch(_ROUTER_CLOCK, return_value=_FIRST_MINUTE):
|
||||
await router.acompletion(**_RPM_LIMITED_REQUEST)
|
||||
with patch(_ROUTER_CLOCK, return_value=_NEXT_MINUTE):
|
||||
next_minute_response: Final = await router.acompletion(**_RPM_LIMITED_REQUEST)
|
||||
|
||||
assert next_minute_response.choices[0].message.content == "response"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue