mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
816947ee61
A streaming request that breaks mid-flight, for example on a mid-stream read
timeout, still bills the provider for the chunks already delivered, yet the proxy
recorded that interrupted request as a zero-spend failure. An earlier revision
logged the recovered partial usage through the success path, which mislabeled a
failed request as a success and produced a misleading spend row
This recovers the partial usage where the failure is actually logged. The
streaming handler assembles the usage from the chunks seen so far and stashes it,
with its cost, on the logging object before firing the failure handlers. The
proxy failure hook lifts that usage and cost onto request_data before the
non-serialisable logging object is popped, and the spend-log writer records the
real partial spend on the failure row instead of a hardcoded zero;
get_logging_payload honors the recovered usage for the token columns and
_failure_handler_helper_fn preserves the recovered cost so the non-DB failure
loggers stay consistent
A request that recovers via a successful fallback is unaffected: the failure hook
only fires when the whole request fails, so the fallback's combined-usage success
row stays the single source of truth and there is no double counting
Resolves LIT-3825
Co-authored-by: veria-ai[bot] <224490171+veria-ai[bot]@users.noreply.github.com>
(cherry picked from commit 4847fa5dd5)
315 lines
11 KiB
Python
315 lines
11 KiB
Python
import datetime as real_datetime
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.proxy._types import ProxyErrorTypes
|
|
from litellm.proxy.utils import ProxyLogging
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.proxy.utils import get_custom_url, join_paths
|
|
|
|
|
|
def test_get_custom_url(monkeypatch):
|
|
monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm")
|
|
custom_url = get_custom_url(request_base_url="http://0.0.0.0:4000", route="ui/")
|
|
assert custom_url == "http://0.0.0.0:4000/litellm/ui/"
|
|
|
|
|
|
def test_proxy_only_error_true_for_llm_route():
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
|
assert proxy_logging_obj._is_proxy_only_llm_api_error(
|
|
original_exception=Exception(),
|
|
error_type=ProxyErrorTypes.auth_error,
|
|
route="/v1/chat/completions",
|
|
)
|
|
|
|
|
|
def test_proxy_only_error_true_for_info_route():
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
|
assert (
|
|
proxy_logging_obj._is_proxy_only_llm_api_error(
|
|
original_exception=Exception(),
|
|
error_type=ProxyErrorTypes.auth_error,
|
|
route="/key/info",
|
|
)
|
|
is True
|
|
)
|
|
|
|
|
|
def test_proxy_only_error_false_for_non_llm_non_info_route():
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
|
assert (
|
|
proxy_logging_obj._is_proxy_only_llm_api_error(
|
|
original_exception=Exception(),
|
|
error_type=ProxyErrorTypes.auth_error,
|
|
route="/key/generate",
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
def test_proxy_only_error_false_for_other_error_type():
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
|
assert (
|
|
proxy_logging_obj._is_proxy_only_llm_api_error(
|
|
original_exception=Exception(),
|
|
error_type=None,
|
|
route="/v1/chat/completions",
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
def test_get_model_group_info_order():
|
|
from litellm import Router
|
|
from litellm.proxy.proxy_server import _get_model_group_info
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "openai/tts-1",
|
|
"litellm_params": {
|
|
"model": "openai/tts-1",
|
|
"api_key": "sk-1234",
|
|
},
|
|
},
|
|
{
|
|
"model_name": "openai/gpt-3.5-turbo",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-3.5-turbo",
|
|
"api_key": "sk-1234",
|
|
},
|
|
},
|
|
]
|
|
)
|
|
model_list = _get_model_group_info(
|
|
llm_router=router,
|
|
all_models_str=["openai/tts-1", "openai/gpt-3.5-turbo"],
|
|
model_group=None,
|
|
)
|
|
|
|
model_groups = [m.model_group for m in model_list]
|
|
assert model_groups == ["openai/tts-1", "openai/gpt-3.5-turbo"]
|
|
|
|
|
|
def test_join_paths_no_duplication():
|
|
"""Test that join_paths doesn't duplicate route when base_path already ends with it"""
|
|
result = join_paths(
|
|
base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path"
|
|
)
|
|
assert result == "http://0.0.0.0:4000/my-custom-path"
|
|
|
|
|
|
def test_join_paths_normal_join():
|
|
"""Test normal path joining"""
|
|
result = join_paths(base_path="http://0.0.0.0:4000", route="/api/v1")
|
|
assert result == "http://0.0.0.0:4000/api/v1"
|
|
|
|
|
|
def test_join_paths_with_trailing_slash():
|
|
"""Test path joining with trailing slash on base_path"""
|
|
result = join_paths(base_path="http://0.0.0.0:4000/", route="api/v1")
|
|
assert result == "http://0.0.0.0:4000/api/v1"
|
|
|
|
|
|
def test_join_paths_empty_base():
|
|
"""Test path joining with empty base_path"""
|
|
result = join_paths(base_path="", route="api/v1")
|
|
assert result == "/api/v1"
|
|
|
|
|
|
def test_join_paths_empty_route():
|
|
"""Test path joining with empty route"""
|
|
result = join_paths(base_path="http://0.0.0.0:4000", route="")
|
|
assert result == "http://0.0.0.0:4000"
|
|
|
|
|
|
def test_join_paths_both_empty():
|
|
"""Test path joining with both empty"""
|
|
result = join_paths(base_path="", route="")
|
|
assert result == "/"
|
|
|
|
|
|
def test_join_paths_nested_path():
|
|
"""Test path joining with nested paths"""
|
|
result = join_paths(base_path="http://0.0.0.0:4000/v1", route="chat/completions")
|
|
assert result == "http://0.0.0.0:4000/v1/chat/completions"
|
|
|
|
|
|
def _patch_today(monkeypatch, year, month, day):
|
|
class PatchedDate(real_datetime.date):
|
|
@classmethod
|
|
def today(cls):
|
|
return real_datetime.date(year, month, day)
|
|
|
|
monkeypatch.setattr("litellm.proxy.utils.date", PatchedDate)
|
|
|
|
|
|
def test_get_projected_spend_over_limit_day_one(monkeypatch):
|
|
from litellm.proxy.utils import _get_projected_spend_over_limit
|
|
|
|
_patch_today(monkeypatch, 2026, 1, 1)
|
|
result = _get_projected_spend_over_limit(100.0, 1.0)
|
|
|
|
assert result is not None
|
|
projected_spend, projected_exceeded_date = result
|
|
assert projected_spend == 3100.0
|
|
assert projected_exceeded_date == real_datetime.date(2026, 1, 1)
|
|
|
|
|
|
def test_get_projected_spend_over_limit_december(monkeypatch):
|
|
from litellm.proxy.utils import _get_projected_spend_over_limit
|
|
|
|
_patch_today(monkeypatch, 2026, 12, 15)
|
|
result = _get_projected_spend_over_limit(100.0, 1.0)
|
|
|
|
assert result is not None
|
|
projected_spend, projected_exceeded_date = result
|
|
assert projected_spend == pytest.approx(214.28571428571428)
|
|
assert projected_exceeded_date == real_datetime.date(2026, 12, 15)
|
|
|
|
|
|
def test_get_projected_spend_over_limit_includes_current_spend(monkeypatch):
|
|
from litellm.proxy.utils import _get_projected_spend_over_limit
|
|
|
|
_patch_today(monkeypatch, 2026, 4, 11)
|
|
result = _get_projected_spend_over_limit(100.0, 200.0)
|
|
|
|
assert result is not None
|
|
projected_spend, projected_exceeded_date = result
|
|
assert projected_spend == 290.0
|
|
assert projected_exceeded_date == real_datetime.date(2026, 4, 21)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# L2: _enrich_http_exception_with_guardrail_context
|
|
# Regression coverage for case 2026-04-10-internal-bedrock-guardrail-streaming-error.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_enrich_http_exception_with_guardrail_context_dict_detail():
|
|
"""L2: dict-detail HTTPException is enriched with guardrail_name and mode."""
|
|
from litellm.proxy.utils import _enrich_http_exception_with_guardrail_context
|
|
|
|
class StubCallback:
|
|
guardrail_name = "bedrock-pii-guard"
|
|
event_hook = "post_call"
|
|
|
|
exc = HTTPException(status_code=400, detail={"error": "Violated guardrail policy"})
|
|
_enrich_http_exception_with_guardrail_context(exc, StubCallback())
|
|
assert exc.detail["guardrail_name"] == "bedrock-pii-guard"
|
|
assert exc.detail["guardrail_mode"] == "post_call"
|
|
|
|
|
|
def test_enrich_http_exception_string_detail_noop():
|
|
"""L2: string-detail HTTPException is not mutated (can't add fields to a str)."""
|
|
from litellm.proxy.utils import _enrich_http_exception_with_guardrail_context
|
|
|
|
class StubCallback:
|
|
guardrail_name = "x"
|
|
event_hook = "pre_call"
|
|
|
|
exc = HTTPException(status_code=400, detail="Content blocked")
|
|
_enrich_http_exception_with_guardrail_context(exc, StubCallback())
|
|
assert exc.detail == "Content blocked"
|
|
|
|
|
|
def test_enrich_http_exception_setdefault_does_not_overwrite():
|
|
"""L2: a guardrail that already populates guardrail_name explicitly wins."""
|
|
from litellm.proxy.utils import _enrich_http_exception_with_guardrail_context
|
|
|
|
class StubCallback:
|
|
guardrail_name = "inferred-name"
|
|
event_hook = "pre_call"
|
|
|
|
exc = HTTPException(
|
|
status_code=400,
|
|
detail={"error": "x", "guardrail_name": "explicit-name"},
|
|
)
|
|
_enrich_http_exception_with_guardrail_context(exc, StubCallback())
|
|
assert exc.detail["guardrail_name"] == "explicit-name"
|
|
|
|
|
|
def test_enrich_http_exception_non_http_exception_noop():
|
|
"""L2: non-HTTPException is left alone and the helper does not raise."""
|
|
from litellm.proxy.utils import _enrich_http_exception_with_guardrail_context
|
|
|
|
class StubCallback:
|
|
guardrail_name = "x"
|
|
event_hook = "pre_call"
|
|
|
|
exc = ValueError("not an HTTPException")
|
|
_enrich_http_exception_with_guardrail_context(exc, StubCallback())
|
|
assert str(exc) == "not an HTTPException"
|
|
|
|
|
|
def test_enrich_http_exception_callback_without_guardrail_name_noop():
|
|
"""L2: callback without guardrail_name attribute leaves detail alone."""
|
|
from litellm.proxy.utils import _enrich_http_exception_with_guardrail_context
|
|
|
|
class StubCallback:
|
|
pass
|
|
|
|
exc = HTTPException(status_code=400, detail={"error": "x"})
|
|
_enrich_http_exception_with_guardrail_context(exc, StubCallback())
|
|
assert exc.detail == {"error": "x"}
|
|
|
|
|
|
class TestPostCallFailureHookLiftsRecoveredPartialSpend:
|
|
"""A stream that broke mid-flight still billed the provider for the chunks
|
|
already delivered. The streaming handler stashes that recovered usage and
|
|
cost on the logging object; post_call_failure_hook must lift them onto
|
|
request_data before the logging object is popped, so the failure-path spend
|
|
callbacks (which run after the pop) record the real partial spend.
|
|
"""
|
|
|
|
async def _run(self, request_data):
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
|
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
|
proxy_logging_obj.alert_types = []
|
|
with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()):
|
|
await proxy_logging_obj.post_call_failure_hook(
|
|
request_data=request_data,
|
|
original_exception=Exception("boom"),
|
|
user_api_key_dict=UserAPIKeyAuth(),
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lifts_recovered_usage_and_cost(self):
|
|
from litellm.types.utils import Usage
|
|
|
|
recovered_usage = Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31)
|
|
logging_obj = MagicMock()
|
|
logging_obj.model_call_details = {
|
|
"combined_usage_object": recovered_usage,
|
|
"response_cost": 3.5e-05,
|
|
}
|
|
request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
|
|
await self._run(request_data)
|
|
|
|
assert request_data["combined_usage_object"] is recovered_usage
|
|
assert request_data["response_cost"] == 3.5e-05
|
|
assert "litellm_logging_obj" not in request_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_recovered_usage_is_noop(self):
|
|
logging_obj = MagicMock()
|
|
logging_obj.model_call_details = {}
|
|
request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
|
|
await self._run(request_data)
|
|
assert "combined_usage_object" not in request_data
|
|
assert "response_cost" not in request_data
|