mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(proxy): per-request header opts a stream into usage.cost without the global include_cost_in_streaming_usage flag (#50574)
* feat(proxy): let one streaming request opt into usage.cost with x-litellm-include-cost-in-streaming-usage The global include_cost_in_streaming_usage flag moved every stream off the fast path. A request can now send x-litellm-include-cost-in-streaming-usage: true to get usage.cost on its usage-carrying stream event (/v1/messages and the Anthropic and OpenAI pass-throughs) while requests without the header keep the fast path. The global flag keeps working unchanged. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: type the new stream cost header tests * fix(types): keep reportUnknownArgumentType at its merge-base count * fix(proxy): the stream cost header only accepts true --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
3c3f266322
commit
07c7229895
9 changed files with 235 additions and 27 deletions
|
|
@ -113,6 +113,27 @@ _Key = TypeVar("_Key") # rebind-ok: a TypeVar cannot be Final
|
|||
_Value = TypeVar("_Value") # rebind-ok: a TypeVar cannot be Final
|
||||
|
||||
|
||||
INCLUDE_COST_IN_STREAMING_USAGE_HEADER: Final = "x-litellm-include-cost-in-streaming-usage"
|
||||
|
||||
|
||||
def proxy_request_headers(container: Mapping[str, object] | None) -> Mapping[str, object] | None:
|
||||
proxy_server_request: Final = as_str_mapping(container.get("proxy_server_request")) if container else None
|
||||
return as_str_mapping(proxy_server_request.get("headers")) if proxy_server_request else None
|
||||
|
||||
|
||||
def stream_cost_injection_requested(request_headers: Mapping[str, object] | None) -> bool:
|
||||
"""The global include_cost_in_streaming_usage flag or the per-request header asks for usage.cost on streams."""
|
||||
if bool(getattr(litellm, "include_cost_in_streaming_usage", False)):
|
||||
return True
|
||||
if not request_headers:
|
||||
return False
|
||||
header_value: Final = next(
|
||||
(value for name, value in request_headers.items() if name.lower() == INCLUDE_COST_IN_STREAMING_USAGE_HEADER),
|
||||
None,
|
||||
)
|
||||
return isinstance(header_value, str) and header_value.strip().lower() == "true"
|
||||
|
||||
|
||||
def without(
|
||||
mapping: Mapping[_Key, _Value], *keys: object
|
||||
) -> dict[_Key, _Value]: # mutable-ok: a fresh copy the caller owns
|
||||
|
|
|
|||
|
|
@ -70,8 +70,10 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
has_attributes,
|
||||
independent_snapshot,
|
||||
is_expected_client_error,
|
||||
proxy_request_headers,
|
||||
redact_nested_match_and_regex_keys,
|
||||
require_attributes,
|
||||
stream_cost_injection_requested,
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
|
|
@ -4815,7 +4817,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# await, response-string materialization, and cost-injection call are
|
||||
# pure overhead on the streaming hot path (the default config).
|
||||
caps: Final = ProxyLogging.callback_capabilities()
|
||||
cost_injection_enabled: Final = bool(getattr(litellm, "include_cost_in_streaming_usage", False))
|
||||
request_headers: Final = proxy_request_headers(
|
||||
request_data # pyright: ignore[reportUnknownArgumentType] # request_data is the untyped proxy request dict
|
||||
)
|
||||
cost_injection_enabled: Final = stream_cost_injection_requested(request_headers)
|
||||
fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled
|
||||
debug_enabled: Final = verbose_proxy_logger.isEnabledFor(logging.DEBUG)
|
||||
stream_completed = False
|
||||
|
|
@ -4862,7 +4867,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
model_name = request_data.get("model", "")
|
||||
chunk = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(
|
||||
chunk, model_name, request_data.get("litellm_logging_obj")
|
||||
chunk, model_name, request_data.get("litellm_logging_obj"), request_headers=request_headers
|
||||
)
|
||||
|
||||
# Set before the yield: an async generator suspends at the yield,
|
||||
|
|
@ -4991,18 +4996,30 @@ class ProxyBaseLLMRequestProcessing:
|
|||
@overload
|
||||
@staticmethod
|
||||
def process_chunk_with_cost_injection(
|
||||
chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None
|
||||
chunk: bytes,
|
||||
model_name: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None = None,
|
||||
*,
|
||||
request_headers: Mapping[str, object] | None = None,
|
||||
) -> bytes: ...
|
||||
|
||||
@overload
|
||||
@staticmethod
|
||||
def process_chunk_with_cost_injection(
|
||||
chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None
|
||||
chunk: object,
|
||||
model_name: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None = None,
|
||||
*,
|
||||
request_headers: Mapping[str, object] | None = None,
|
||||
) -> object: ...
|
||||
|
||||
@staticmethod
|
||||
def process_chunk_with_cost_injection(
|
||||
chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None
|
||||
chunk: object,
|
||||
model_name: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None = None,
|
||||
*,
|
||||
request_headers: Mapping[str, object] | None = None,
|
||||
) -> object:
|
||||
"""
|
||||
Process a streaming chunk and inject cost information if enabled.
|
||||
|
|
@ -5011,11 +5028,12 @@ class ProxyBaseLLMRequestProcessing:
|
|||
chunk: The streaming chunk (dict, str, bytes, or bytearray)
|
||||
model_name: Model name for cost calculation
|
||||
litellm_logging_obj: The call's logging object, used for pricing
|
||||
request_headers: The proxy request headers, read for the per-request opt-in header
|
||||
|
||||
Returns:
|
||||
The processed chunk with cost information injected if applicable
|
||||
"""
|
||||
if not getattr(litellm, "include_cost_in_streaming_usage", False):
|
||||
if not stream_cost_injection_requested(request_headers):
|
||||
return chunk
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -6,10 +6,14 @@ from typing import Final, Protocol
|
|||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_str_mapping,
|
||||
bind_budget_reservation_to_callbacks,
|
||||
proxy_request_headers,
|
||||
stream_cost_injection_requested,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception
|
||||
|
|
@ -194,12 +198,11 @@ class PassThroughStreamingHandler:
|
|||
litellm_logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
# Resolve once per stream rather than re-reading the global +
|
||||
# re-branching on every chunk. ``include_cost_in_streaming_usage`` is
|
||||
# set at config load and stable for the process, matching how the
|
||||
# proxy-level streaming fast path resolves it.
|
||||
request_headers: Final = proxy_request_headers(
|
||||
as_str_mapping(litellm_logging_obj.model_call_details.get("litellm_params"))
|
||||
)
|
||||
cost_injection_active: Final = (
|
||||
bool(getattr(litellm, "include_cost_in_streaming_usage", False))
|
||||
stream_cost_injection_requested(request_headers)
|
||||
and bool(model_name)
|
||||
and (
|
||||
endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI)
|
||||
|
|
@ -242,7 +245,7 @@ class PassThroughStreamingHandler:
|
|||
complete_frames, pending = split_complete_sse_frames(pending + chunk)
|
||||
if complete_frames:
|
||||
yield ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(
|
||||
complete_frames, resolved_model_name, litellm_logging_obj
|
||||
complete_frames, resolved_model_name, litellm_logging_obj, request_headers=request_headers
|
||||
)
|
||||
if pending:
|
||||
yield pending
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ class Feature(str, Enum):
|
|||
CACHE_CONTROL_INJECTION = "cache_control_injection"
|
||||
"""Prompt caching injection that Python applies to messages and system."""
|
||||
STREAM_COST_INJECTION = "stream_cost_injection"
|
||||
"""litellm.include_cost_in_streaming_usage, where Python adds cost to streamed usage."""
|
||||
"""litellm.include_cost_in_streaming_usage or the x-litellm-include-cost-in-streaming-usage header, where Python adds cost to streamed usage."""
|
||||
THINKING_SIGNATURE_RETRY = "thinking_signature_retry"
|
||||
"""Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature."""
|
||||
TOOL_CUSTOM_FIELD = "tool_custom_field"
|
||||
|
|
|
|||
|
|
@ -6,7 +6,11 @@ from typing import Final
|
|||
|
||||
import litellm
|
||||
from litellm.anthropic_beta_headers_manager import is_foundry_api_base
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
drops_unsupported_params,
|
||||
proxy_request_headers,
|
||||
stream_cost_injection_requested,
|
||||
)
|
||||
from litellm.llms.opencode.common_utils import is_opencode_api_base
|
||||
from litellm.rust_bridge.detectors import api_base, overridden, present, streaming
|
||||
from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str
|
||||
|
|
@ -64,8 +68,8 @@ def strict_tools_retry(request: NativeCall) -> bool:
|
|||
|
||||
|
||||
def stream_cost_injection(request: NativeCall) -> bool:
|
||||
"""Python adds response cost to streaming usage when the global setting is truthy."""
|
||||
return streaming(request) and bool(litellm.include_cost_in_streaming_usage)
|
||||
"""Python adds response cost to streaming usage when the global setting or the request header asks for it."""
|
||||
return streaming(request) and stream_cost_injection_requested(proxy_request_headers(request.resolved))
|
||||
|
||||
|
||||
def pre_request_hook(_request: NativeCall) -> bool:
|
||||
|
|
|
|||
|
|
@ -80,7 +80,9 @@ def _stream(message_id: str) -> bytes:
|
|||
)
|
||||
|
||||
|
||||
def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path:
|
||||
def _proxy_config(
|
||||
directory: Path, model: str, upstream_url: str, litellm_settings: dict[str, JsonValue] | None = None
|
||||
) -> Path:
|
||||
config: Final = directory / "streamed_usage_cost_config.yaml"
|
||||
config.write_text(
|
||||
json.dumps(
|
||||
|
|
@ -109,7 +111,9 @@ def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path:
|
|||
"proxy_batch_write_at": 1,
|
||||
"proxy_batch_polling_interval": 1,
|
||||
},
|
||||
"litellm_settings": {"include_cost_in_streaming_usage": True},
|
||||
"litellm_settings": {"include_cost_in_streaming_usage": True}
|
||||
if litellm_settings is None
|
||||
else litellm_settings,
|
||||
"router_settings": {"disable_cooldowns": True},
|
||||
}
|
||||
)
|
||||
|
|
@ -176,3 +180,49 @@ def test_bedrock_messages_stream_usage_cost_matches_recorded_spend_with_custom_c
|
|||
"cost": pytest.approx(recorded_spend),
|
||||
}, (streamed_usage, rows, response.text)
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_messages_stream_header_opts_into_usage_cost_while_global_flag_is_off(gateway: Gateway, tmp_path: Path) -> None:
|
||||
model: Final = f"{BEDROCK_MODEL}-{uuid.uuid4().hex}"
|
||||
message_ids: Final = {"plain": f"msg_{uuid.uuid4().hex}", "opted in": f"msg_{uuid.uuid4().hex}"}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.target == f"/model/{BEDROCK_MODEL}/invoke-with-response-stream", request.target
|
||||
prompt: Final = json.loads(request.body)["messages"][0]["content"]
|
||||
return Reply(content_type="application/vnd.amazon.eventstream", chunks=(_stream(message_ids[prompt]),))
|
||||
|
||||
def message_delta_usage(prompt: str, headers: dict[str, str]) -> dict[str, float]:
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": OUTPUT_TOKENS,
|
||||
"stream": True,
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
message_delta: Final = next(event for event in _data_events(response.text) if event["type"] == "message_delta")
|
||||
return STREAMED_USAGE.validate_python(message_delta["usage"])
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
config: Final = _proxy_config(tmp_path, model, wire.url, litellm_settings={})
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate:
|
||||
plain_usage: Final = message_delta_usage("plain", {})
|
||||
opted_in_usage: Final = message_delta_usage(
|
||||
"opted in", {"x-litellm-include-cost-in-streaming-usage": "true"}
|
||||
)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (message_ids["opted in"],)
|
||||
),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert "cost" not in plain_usage, plain_usage
|
||||
assert opted_in_usage["cost"] == pytest.approx(float(str(rows[0]["spend"]))), (opted_in_usage, rows)
|
||||
assert opted_in_usage["cost"] == pytest.approx(EXPECTED_SPEND), opted_in_usage
|
||||
assert len(wire.drain()) == 2
|
||||
|
|
|
|||
|
|
@ -445,6 +445,47 @@ async def test_chunk_processor_injects_cost_into_openai_passthrough_usage_frame(
|
|||
assert final_payload["usage"]["total_tokens"] == 15
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("header_value", "cost_expected"), (("true", True), ("TRUE", True), ("1", False), ("false", False), (None, False))
|
||||
)
|
||||
async def test_chunk_processor_honors_the_opt_in_header_while_global_flag_is_off(
|
||||
monkeypatch: pytest.MonkeyPatch, header_value: str | None, cost_expected: bool
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
|
||||
chunks: Final = _openai_passthrough_stream_chunks()
|
||||
logging_obj: Final = _unarmed_logging_obj(model="gpt-4o-mini")
|
||||
headers: Final = {} if header_value is None else {"x-litellm-include-cost-in-streaming-usage": header_value}
|
||||
logging_obj.model_call_details = {"litellm_params": {"proxy_server_request": {"headers": headers}}}
|
||||
|
||||
received: Final = [
|
||||
chunk
|
||||
async for chunk in PassThroughStreamingHandler.chunk_processor(
|
||||
response=_make_streaming_response(chunks),
|
||||
request_body={"model": "gpt-4o-mini", "stream": True},
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=EndpointType.OPENAI,
|
||||
start_time=datetime(2026, 1, 1),
|
||||
passthrough_success_handler_obj=MagicMock(),
|
||||
url_route="/openai/v1/chat/completions",
|
||||
route_streaming_logging=AsyncMock(),
|
||||
)
|
||||
]
|
||||
await asyncio.sleep(0)
|
||||
|
||||
usage_payloads: Final = [
|
||||
json.loads(chunk.decode("utf-8").split("data:", 1)[1].strip())
|
||||
for chunk in received
|
||||
if b'"total_tokens"' in chunk
|
||||
]
|
||||
assert len(usage_payloads) == 1
|
||||
assert ("cost" in usage_payloads[0]["usage"]) is cost_expected
|
||||
if cost_expected:
|
||||
assert usage_payloads[0]["usage"]["cost"] > 0
|
||||
else:
|
||||
assert received == chunks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_injects_cost_into_usage_frame_fragmented_across_chunks(monkeypatch):
|
||||
"""Regression: an SSE usage frame split across transport chunks must still get
|
||||
|
|
@ -826,7 +867,8 @@ async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logg
|
|||
|
||||
logging_obj.carry_failed_attempt_usage()
|
||||
retry: Final = _anthropic_messages_attempt(
|
||||
logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail()))
|
||||
logging_obj,
|
||||
_anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())),
|
||||
)
|
||||
async for _ in retry:
|
||||
pass
|
||||
|
|
@ -858,6 +900,7 @@ async def test_every_dropped_attempt_logs_its_own_failure() -> None:
|
|||
stashed: Final = logging_obj.model_call_details.get("combined_usage_object")
|
||||
assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9
|
||||
|
||||
|
||||
def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes:
|
||||
return _anthropic_sse(
|
||||
"message_start",
|
||||
|
|
@ -890,7 +933,11 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]:
|
|||
_anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
_anthropic_sse(
|
||||
"message_delta",
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 3}},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 3},
|
||||
},
|
||||
),
|
||||
_anthropic_sse("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
|
@ -923,7 +970,9 @@ def _anthropic_stream_of(frames: tuple[bytes, ...]):
|
|||
"frames, billed",
|
||||
[
|
||||
pytest.param(
|
||||
(_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), True, id="error-after-message-start"
|
||||
(_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()),
|
||||
True,
|
||||
id="error-after-message-start",
|
||||
),
|
||||
pytest.param((_anthropic_overloaded_error_frame(),), False, id="error-frame-only"),
|
||||
],
|
||||
|
|
@ -1102,7 +1151,8 @@ async def test_error_frame_then_exhausted_retry_bills_both_attempts() -> None:
|
|||
recorder: Final = _EventRecorder()
|
||||
logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-error-frame-terminal-billing")
|
||||
superseded: Final = _anthropic_messages_attempt(
|
||||
logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame()))
|
||||
logging_obj,
|
||||
_anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())),
|
||||
)
|
||||
async for chunk in superseded:
|
||||
if b"event: error" in chunk:
|
||||
|
|
@ -1168,6 +1218,8 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr
|
|||
success_payload = recorder.success_kwargs[0]["standard_logging_object"]
|
||||
assert success_payload["status"] == "success"
|
||||
assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3)
|
||||
|
||||
|
||||
def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes:
|
||||
payload = {
|
||||
"candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}],
|
||||
|
|
|
|||
|
|
@ -5284,9 +5284,7 @@ class TestHandleLLMApiExceptionProviderOverload:
|
|||
assert proxy_exc.type == "internal_server_error"
|
||||
assert proxy_exc.to_dict()["code"] == "500"
|
||||
|
||||
async def test_openai_route_answers_wrapped_provider_overload_as_503(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
async def test_openai_route_answers_wrapped_provider_overload_as_503(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
original: Final = InternalServerError(
|
||||
message="Overloaded",
|
||||
llm_provider="anthropic",
|
||||
|
|
@ -5517,6 +5515,43 @@ class TestAsyncStreamingDataGeneratorFastPath:
|
|||
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
async def test_opt_in_header_injects_cost_while_global_flag_is_off(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
usage_frame: Final = (
|
||||
b'data: {"object":"chat.completion.chunk","choices":[],'
|
||||
b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n'
|
||||
)
|
||||
|
||||
async def frames() -> AsyncGenerator[bytes, None]:
|
||||
yield usage_frame
|
||||
|
||||
async def stream(headers: dict[str, str]) -> list[object]:
|
||||
return [
|
||||
c
|
||||
async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=frames(),
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
request_data={"model": "gpt-4o-mini", "proxy_server_request": {"headers": headers}},
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()),
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
|
||||
serialize_error=lambda e: "data: error\n\n",
|
||||
)
|
||||
]
|
||||
|
||||
plain: Final = await stream({})
|
||||
opted_in: Final = await stream({"x-litellm-include-cost-in-streaming-usage": "true"})
|
||||
declined: Final = await stream({"x-litellm-include-cost-in-streaming-usage": "false"})
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
assert plain == [usage_frame]
|
||||
assert declined == [usage_frame]
|
||||
assert isinstance(opted_in[0], bytes)
|
||||
payload: Final = json.loads(opted_in[0].decode("utf-8").split("data:", 1)[1].strip())
|
||||
assert payload["usage"]["cost"] > 0
|
||||
assert payload["usage"]["total_tokens"] == 15
|
||||
|
||||
|
||||
class TestDisconnectGatherCleanup:
|
||||
def _disconnect_request(self) -> Request:
|
||||
|
|
@ -10823,6 +10858,23 @@ class TestProcessChunkWithCostInjection:
|
|||
payload = json.loads(result.decode("utf-8").split("data:", 1)[1].strip())
|
||||
assert payload["usage"]["cost"] > 0
|
||||
|
||||
def test_request_header_opts_a_frame_into_cost_while_global_flag_is_off(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
|
||||
chunk: Final = (
|
||||
b'data: {"object":"chat.completion.chunk","choices":[],'
|
||||
b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n'
|
||||
)
|
||||
|
||||
assert ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk
|
||||
injected: Final = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(
|
||||
chunk, "gpt-4o-mini", request_headers={"X-LiteLLM-Include-Cost-In-Streaming-Usage": "True"}
|
||||
)
|
||||
|
||||
payload: Final = json.loads(injected.decode("utf-8").split("data:", 1)[1].strip())
|
||||
assert payload["usage"]["cost"] > 0
|
||||
|
||||
def test_chunk_ending_in_partial_frame_passes_through_byte_identical(self, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
|
||||
chunk = (
|
||||
|
|
|
|||
|
|
@ -352,6 +352,14 @@ def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) -
|
|||
assert not messages_detectors.stream_cost_injection(call())
|
||||
|
||||
|
||||
def test_stream_cost_injection_reads_the_request_header_the_proxy_reads(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
|
||||
opted_in: Final = {"proxy_server_request": {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}}
|
||||
assert messages_detectors.stream_cost_injection(call(stream=True, **opted_in))
|
||||
assert not messages_detectors.stream_cost_injection(call(**opted_in))
|
||||
assert not messages_detectors.stream_cost_injection(call(stream=True, proxy_server_request={"headers": {}}))
|
||||
|
||||
|
||||
def test_foundry_beta_row_reads_the_api_base_the_anthropic_handler_reads(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue