scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

fix(vllm): write spend log rows for /vllm passthrough calls to router models (#45733)

Browse source 
* fix(vllm): write spend log rows for /vllm passthrough calls to router models

The vLLM passthrough config never built a logged response for relayed calls, so stream and non-stream calls through /vllm with a router model wrote no spend row and never charged the key. vLLM now shares the OpenAI-shaped relay logging that Azure already had, moved into an openai_like helper both configs call.

* ci: run the vllm passthrough unit tests in the llms-providers shard

* fix(proxy): log /vllm text completions and ask router models for stream usage

A chat stream relayed to a router model without stream_options now gets
include_usage added on the way out, and the usage chunk vLLM appends is
removed before the client sees it, so the spend row carries vLLM's own
token counts instead of a local estimate. /vllm/completions calls now
write a row too, stream and non-stream.

* fix(vllm): keep usage, fingerprints, and disconnect rows correct on /vllm streams

Inject and strip stream usage on /vllm/completions too, hold every finish
frame for the fingerprint when n > 1, pass a forced-usage server's stream
through untouched, stop turning continuous_usage_stats alone into per-chunk
usage, close the upstream when the client hangs up so the row is written,
and assemble text choices with one sort instead of a scan per choice

* fix(vllm): keep n>1 stream order when holding finish chunks for the fingerprint

A held finish chunk now goes out as soon as any other chunk arrives, so a
choice that finishes while another still streams keeps its place. Only the
finish chunks right before the usage chunk get vLLM's system_fingerprint.

* fix(vllm): keep caller stream_options and log both success handlers

Forward a stream_options the proxy cannot read, or a non-boolean
include_usage, to vLLM as the caller sent it. Only a false
always_include_stream_usage changes the /vllm relay, so clients get the
same stream bytes under every setting.

Record the passthrough origin when logging starts so the sync and async
success handlers both parse the relayed body, hand callbacks a parsed
text completion, and read relayed messages and prompt separately so a
stray field no longer drops the local count.

Add the audit's integration tests for the relay, the settings, and the
Azure spend row.

* test(vllm): expect the caller's stream under always_include_stream_usage true

* fix(vllm): relay long SSE frames in linear time

* test(vllm): time the relay's usage and finish_reason prefilter on one large frame

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-14 12:40:48 -07:00 • committed by GitHub
parent 039dea2642
commit 2d4472efc1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 2673 additions and 115 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -350,6 +350,7 @@ jobs:
tests/unit/llms/v0
tests/unit/llms/valkey
tests/unit/llms/vercel_ai_gateway
tests/unit/llms/vllm
tests/unit/llms/volcengine
tests/unit/llms/voyage
tests/unit/llms/wandb

13
litellm/litellm_core_utils/litellm_logging.py
View file

@ -651,6 +651,10 @@ class Logging(LiteLLMLoggingBaseClass):
self.stream = stream
self.start_time = start_time # log the call start time
self.call_type = call_type
self.relays_llm_passthrough: Final = call_type in (
CallTypes.llm_passthrough_route.value,
CallTypes.allm_passthrough_route.value,
)
self.litellm_call_id = litellm_call_id
self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
@ -2388,10 +2392,7 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier=ws_service_tier,
)
elif (
self.call_type == CallTypes.llm_passthrough_route.value
or self.call_type == CallTypes.allm_passthrough_route.value
) and isinstance(result, Response):
elif self.relays_llm_passthrough and isinstance(result, Response):
from litellm.utils import ProviderConfigManager
provider_config: Final = ProviderConfigManager.get_provider_passthrough_config(
@ -2588,7 +2589,9 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result: Final = self.normalize_logging_result(result=result)
if isinstance(result, Response) and isinstance(logging_result, (ModelResponse, EmbeddingResponse)):
if isinstance(result, Response) and isinstance(
logging_result, (ModelResponse, EmbeddingResponse, TextCompletionResponse)
):
result = logging_result
if standard_logging_object is None and result is not None and self.stream is not True:

86
litellm/llms/azure/passthrough/transformation.py
View file

@ -5,22 +5,18 @@ from typing import TYPE_CHECKING, Final, Optional
import httpx
from httpx import Response
from pydantic import ValidationError
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.base_llm.passthrough.transformation import (
BasePassthroughConfig,
RelayShape,
logged_relay_shape,
replace_path_segment,
strip_leading_model_segment,
)
from litellm.llms.openai_like.passthrough.transformation import logged_openai_like_response, logged_openai_like_stream
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse
if TYPE_CHECKING:
from httpx import URL
@ -28,42 +24,6 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
class RelayedChatRequest(LiteLLMBaseModel):
messages: Sequence[Mapping[str, object]] | None = None
class RelayedCallDetails(LiteLLMBaseModel):
request_data: RelayedChatRequest | None = None
def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None:
try:
details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details)
except ValidationError:
return None
return details.request_data.messages if details.request_data else None
RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate)
OPENAI_RELAY_SHAPES: Final = (
RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate),
RESPONSES_RELAY_SHAPE,
RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate),
)
def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None:
"""A streaming logging object assembles the logged response from the terminal event, not from its body."""
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks)
if terminal_event is None:
return None
logging_obj.call_type = RESPONSES_RELAY_SHAPE.call_type.value
return terminal_event
AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(?<![^/])openai/deployments/([^/]+)")
AZURE_BODY_MODEL_INFERENCE_ENDPOINTS: Final = frozenset(
{"responses", "chat/completions", "completions", "embeddings", "images/generations", "audio/speech"}
@ -179,33 +139,11 @@ class AzurePassthroughConfig(BasePassthroughConfig):
model: str,
custom_llm_provider: str,
httpx_response: Response,
request_data: dict,
request_data: Mapping[str, object],
logging_obj: Logging,
endpoint: str,
) -> Optional["LoggedRelayResponse"]:
from litellm import encoding
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.types.utils import ModelResponse
if "chat/completions" not in endpoint:
return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint)
openai_chat_config: Final = OpenAIGPTConfig()
litellm_model_response: Final[ModelResponse] = openai_chat_config.transform_response(
model=model,
messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}],
raw_response=httpx_response,
model_response=ModelResponse(),
logging_obj=logging_obj,
optional_params={},
litellm_params={},
api_key="",
request_data=request_data,
encoding=encoding,
)
return litellm_model_response
return logged_openai_like_response(model, httpx_response, request_data, logging_obj, endpoint)
def handle_logging_collected_chunks(
self,
@ -215,18 +153,4 @@ class AzurePassthroughConfig(BasePassthroughConfig):
custom_llm_provider: str,
endpoint: str,
) -> Optional["LoggedRelayResponse"]:
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix):
return logged_responses_stream(all_chunks, litellm_logging_obj)
if "chat/completions" not in endpoint:
return None
return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
messages=_relayed_messages(litellm_logging_obj),
)
return logged_openai_like_stream(all_chunks, litellm_logging_obj, model, endpoint)

249
litellm/llms/openai_like/passthrough/transformation.py Normal file
View file

@ -0,0 +1,249 @@
import itertools
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeAlias
from httpx import Response
from pydantic import ConfigDict, ValidationError
import litellm.litellm_core_utils.token_counter as token_counting
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.passthrough.transformation import RelayShape, logged_relay_shape
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent
from litellm.types.utils import (
CallTypes,
EmbeddingResponse,
ImageResponse,
ModelResponse,
TextChoices,
TextCompletionResponse,
Usage,
)
if TYPE_CHECKING:
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
RelayedPrompt: TypeAlias = str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]]
def _count_text_tokens(model: str, text: str, count_response_tokens: bool = False) -> int:
return token_counting.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's custom_tokenizer param is an untyped dict
model=model, text=text, count_response_tokens=count_response_tokens
)
class RelayedChatRequest(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
messages: Sequence[Mapping[str, object]] | None = None
class RelayedTextRequest(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
prompt: RelayedPrompt | None = None
class RelayedChatCall(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
request_data: RelayedChatRequest | None = None
class RelayedTextCall(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
request_data: RelayedTextRequest | None = None
def relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None:
try:
request: Final = RelayedChatCall.model_validate(litellm_logging_obj.model_call_details).request_data
except ValidationError:
return None
return request.messages if request else None
def _relayed_prompt(litellm_logging_obj: Logging) -> RelayedPrompt | None:
try:
request: Final = RelayedTextCall.model_validate(litellm_logging_obj.model_call_details).request_data
except ValidationError:
return None
return request.prompt if request else None
class RelayedTextChoice(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
index: int = 0
text: str = ""
finish_reason: str | None = None
class RelayedTextCompletionChunk(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
id: str
created: int
model: str
choices: Sequence[RelayedTextChoice] = ()
usage: Usage | None = None
RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate)
TEXT_COMPLETION_RELAY_SHAPE: Final = RelayShape(
"/completions", CallTypes.atext_completion, TextCompletionResponse.model_validate
)
OPENAI_RELAY_SHAPES: Final = (
RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate),
RESPONSES_RELAY_SHAPE,
RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate),
)
def relays_text_completion(endpoint: str) -> bool:
relayed_path: Final = f"/{endpoint.strip('/')}"
return relayed_path.endswith(TEXT_COMPLETION_RELAY_SHAPE.path_suffix) and not relayed_path.endswith(
"/chat/completions"
)
def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None:
"""A streaming logging object assembles the logged response from the terminal event, not from its body."""
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks)
if terminal_event is None:
return None
logging_obj.call_type = RESPONSES_RELAY_SHAPE.call_type.value
return terminal_event
def _relayed_text_chunk(line: str) -> RelayedTextCompletionChunk | None:
payload: Final = BaseModelResponseIterator.string_to_dict_parser(str_line=line)
if not payload:
return None
try:
return RelayedTextCompletionChunk.model_validate(payload)
except ValidationError:
return None
def _assembled_text_choice(index: int, own: Sequence[RelayedTextChoice]) -> RelayedTextChoice:
return RelayedTextChoice(
index=index,
text="".join(delta.text for delta in own),
finish_reason=next((delta.finish_reason for delta in reversed(own) if delta.finish_reason), None),
)
def _delta_index(delta: RelayedTextChoice) -> int:
return delta.index
def _prompt_part_tokens(model: str, part: str | int | Sequence[int]) -> int:
match part:
case str():
return _count_text_tokens(model=model, text=part)
case int():
return 1
case _:
return len(part)
def _prompt_tokens(model: str, prompt: RelayedPrompt | None) -> int:
match prompt:
case None:
return 0
case str():
return _count_text_tokens(model=model, text=prompt)
case _:
return sum(_prompt_part_tokens(model, part) for part in prompt)
def _counted_text_usage(model: str, prompt: RelayedPrompt | None, choices: Sequence[RelayedTextChoice]) -> Usage:
prompt_tokens: Final = _prompt_tokens(model, prompt)
completion_tokens: Final = sum(
_count_text_tokens(model=model, text=choice.text, count_response_tokens=True) for choice in choices
)
return Usage(
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens
)
def logged_text_completion_response(
httpx_response: Response, logging_obj: Logging, endpoint: str
) -> "LoggedRelayResponse | None":
return logged_relay_shape((TEXT_COMPLETION_RELAY_SHAPE,), httpx_response, logging_obj, endpoint)
def logged_text_completion_stream(
all_chunks: Sequence[str], logging_obj: Logging, model: str
) -> TextCompletionResponse | None:
chunks: Final = tuple(chunk for chunk in map(_relayed_text_chunk, all_chunks) if chunk is not None)
head: Final = next((chunk for chunk in chunks if chunk.choices), None)
if head is None:
return None
deltas: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
choices: Final = tuple(
_assembled_text_choice(index, tuple(own))
for index, own in itertools.groupby(sorted(deltas, key=_delta_index), key=_delta_index)
)
reported_usage: Final = next((chunk.usage for chunk in reversed(chunks) if chunk.usage is not None), None)
prompt: Final = _relayed_prompt(logging_obj)
logging_obj.call_type = TEXT_COMPLETION_RELAY_SHAPE.call_type.value # rebind-ok: the relay's logging object starts as the generic passthrough call and learns its real call type from the relayed shape
return TextCompletionResponse(
id=head.id,
created=head.created,
model=head.model,
choices=[
TextChoices(index=choice.index, text=choice.text, finish_reason=choice.finish_reason) for choice in choices
],
usage=(reported_usage if reported_usage is not None else _counted_text_usage(model, prompt, choices)),
)
def logged_openai_like_response(
model: str, httpx_response: Response, request_data: Mapping[str, object], logging_obj: Logging, endpoint: str
) -> "LoggedRelayResponse | None":
from litellm import encoding
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
if "chat/completions" not in endpoint:
return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint)
return OpenAIGPTConfig().transform_response(
model=model,
messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}],
raw_response=httpx_response,
model_response=ModelResponse(),
logging_obj=logging_obj,
optional_params={},
litellm_params={},
api_key="",
request_data=dict(request_data),
encoding=encoding,
)
def logged_openai_like_stream(
all_chunks: Sequence[str], logging_obj: Logging, model: str, endpoint: str
) -> "LoggedRelayResponse | None":
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
relayed_path: Final = f"/{endpoint.strip('/')}"
if relayed_path.endswith(RESPONSES_RELAY_SHAPE.path_suffix):
return logged_responses_stream(all_chunks, logging_obj)
if "chat/completions" in endpoint:
return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser
all_chunks=all_chunks,
litellm_logging_obj=logging_obj,
model=model,
messages=relayed_messages(logging_obj),
)
return None

38
litellm/llms/vllm/passthrough/transformation.py
View file

@ -1,12 +1,25 @@
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
from httpx import Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.llms.openai_like.passthrough.transformation import (
logged_openai_like_response,
logged_openai_like_stream,
logged_text_completion_response,
logged_text_completion_stream,
relays_text_completion,
)
from ..common_utils import VLLMModelInfo
if TYPE_CHECKING:
from httpx import URL
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig):
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
@ -30,3 +43,28 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig):
self.format_url(endpoint, base_target_url, request_query_params),
base_target_url,
)
def logging_non_streaming_response(
self,
model: str,
custom_llm_provider: str,
httpx_response: Response,
request_data: Mapping[str, object],
logging_obj: Logging,
endpoint: str,
) -> "LoggedRelayResponse | None":
if relays_text_completion(endpoint):
return logged_text_completion_response(httpx_response, logging_obj, endpoint)
return logged_openai_like_response(model, httpx_response, request_data, logging_obj, endpoint)
def handle_logging_collected_chunks(
self,
all_chunks: Sequence[str],
litellm_logging_obj: Logging,
model: str,
custom_llm_provider: str,
endpoint: str,
) -> "LoggedRelayResponse | None":
if relays_text_completion(endpoint):
return logged_text_completion_stream(all_chunks, litellm_logging_obj, model)
return logged_openai_like_stream(all_chunks, litellm_logging_obj, model, endpoint)

8
litellm/proxy/common_request_processing.py
View file

@ -513,7 +513,7 @@ def _litellm_model_supports_stream_options(litellm_model: str) -> bool:
return supported_params is not None and "stream_options" in supported_params
def _model_deployments_support_stream_options(
def model_deployments_support_stream_options(
model: object,
llm_router: Router | None,
team_id: str | None,
@ -527,7 +527,7 @@ def _model_deployments_support_stream_options(
return all(_litellm_model_supports_stream_options(m) for m in candidate_models)
def _stream_usage_tracking_updates(
def stream_usage_tracking_updates(
data: Mapping[str, object],
general_settings: Mapping[str, object],
route_type: str,
@ -2277,11 +2277,11 @@ class ProxyBaseLLMRequestProcessing:
### AUTO STREAM USAGE TRACKING ###
self.data.update(
_stream_usage_tracking_updates(
stream_usage_tracking_updates(
data=self.data,
general_settings=general_settings,
route_type=route_type,
supports_stream_options=lambda: _model_deployments_support_stream_options(
supports_stream_options=lambda: model_deployments_support_stream_options(
model=self.data.get("model"),
llm_router=llm_router,
team_id=user_api_key_dict.team_id,

210
litellm/proxy/pass_through_endpoints/injected_stream_usage.py Normal file
View file

@ -0,0 +1,210 @@
import json
import re
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass, replace
from itertools import islice
from types import MappingProxyType
from typing import Final
from pydantic import ConfigDict, TypeAdapter, ValidationError
from litellm.proxy.common_request_processing import stream_usage_tracking_updates
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import CallTypes
_USAGE_FRAME_ENDPOINTS: Final = ("chat/completions", "completions")
_CALLER_ONLY_STREAM_OPTIONS: Final = frozenset({"continuous_usage_stats"})
_SSE_FRAME_END: Final = re.compile(rb"\r\n\r\n|\n\n|\r\r")
_SSE_FRAME_END_LOOKBACK: Final = 3
_USAGE_OR_FINISH: Final = re.compile(rb'"usage"\s*:\s*\{|"finish_reason"\s*:\s*"')
_SSE_DATA_FIELD: Final = b"data:"
_JSON_OBJECT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_INJECTION_OFF: Final[Mapping[str, object]] = MappingProxyType({"always_include_stream_usage": False})
_DEFAULT_INJECTION: Final[Mapping[str, object]] = MappingProxyType({})
class _RelayedChoice(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
finish_reason: str | None = None
class _RelayedChunk(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
choices: Sequence[_RelayedChoice]
usage: Mapping[str, object] | None = None
system_fingerprint: str | None = None
@dataclass(frozen=True, slots=True)
class StreamUsageRelay:
body: Mapping[str, object]
strips_usage: bool
@dataclass(frozen=True, slots=True)
class _Stacked:
item: bytes
below: "_Stacked | None"
@dataclass(frozen=True, slots=True)
class _RelayState:
newest_held: _Stacked | None = None
passes_usage_through: bool = False
def _newest_first(top: _Stacked | None) -> Iterator[bytes]:
while top is not None:
yield top.item
top = top.below # rebind-ok: walks the stack back from the newest item
def _oldest_first(top: _Stacked | None) -> tuple[bytes, ...]:
return tuple(reversed(tuple(_newest_first(top))))
def _held_frames(state: _RelayState) -> tuple[bytes, ...]:
return _oldest_first(state.newest_held)
def _ends_a_frame(unterminated: _Stacked | None, piece: bytes) -> bool:
lookback: Final = b"".join(reversed(tuple(islice(_newest_first(unterminated), _SSE_FRAME_END_LOOKBACK))))
return _SSE_FRAME_END.search(lookback[-_SSE_FRAME_END_LOOKBACK:] + piece) is not None
def _relays_usage_frames(endpoint: str) -> bool:
path: Final = endpoint.strip("/")
return any(path == suffix or path.endswith(f"/{suffix}") for suffix in _USAGE_FRAME_ENDPOINTS)
def _forwarded_stream_options(stream_options: object, strips_usage: bool) -> object:
if not strips_usage or not isinstance(stream_options, Mapping):
return stream_options
caller_options: Final = _JSON_OBJECT_ADAPTER.validate_python(stream_options)
return {key: value for key, value in caller_options.items() if key not in _CALLER_ONLY_STREAM_OPTIONS}
def _relay_settings(general_settings: Mapping[str, object]) -> Mapping[str, object]:
return _INJECTION_OFF if general_settings.get("always_include_stream_usage") is False else _DEFAULT_INJECTION
def _forwards_stream_options_as_sent(request_body: Mapping[str, object]) -> bool:
stream_options: Final = request_body.get("stream_options")
if stream_options is None:
return False
if not isinstance(stream_options, Mapping):
return True
include_usage: Final = _JSON_OBJECT_ADAPTER.validate_python(stream_options).get("include_usage", False)
return not isinstance(include_usage, bool)
def stream_usage_relay(
endpoint: str,
request_body: Mapping[str, object],
general_settings: Mapping[str, object],
supports_stream_options: Callable[[], bool],
) -> StreamUsageRelay:
if not _relays_usage_frames(endpoint) or _forwards_stream_options_as_sent(request_body):
return StreamUsageRelay(body=request_body, strips_usage=False)
updates: Final = stream_usage_tracking_updates(
data=request_body,
general_settings=_relay_settings(general_settings),
route_type=CallTypes.acompletion.value,
supports_stream_options=supports_stream_options,
)
strips_usage: Final = updates.get("_litellm_strip_stream_usage") is True
stream_options: Final = updates.get("stream_options")
return StreamUsageRelay(
body=(
request_body
if stream_options is None
else {**request_body, "stream_options": _forwarded_stream_options(stream_options, strips_usage)}
),
strips_usage=strips_usage,
)
def _frame_payload(frame: bytes) -> bytes:
return b"\n".join(
line[len(_SSE_DATA_FIELD) :].strip() for line in frame.splitlines() if line.startswith(_SSE_DATA_FIELD)
)
def _relayed_chunk(frame: bytes) -> _RelayedChunk | None:
if _USAGE_OR_FINISH.search(frame) is None:
return None
try:
return _RelayedChunk.model_validate_json(_frame_payload(frame))
except ValidationError:
return None
def _with_fingerprint(frame: bytes, fingerprint: str | None) -> bytes:
if fingerprint is None:
return frame
payload: Final = _JSON_OBJECT_ADAPTER.validate_json(_frame_payload(frame))
terminator: Final = frame[len(frame.rstrip(b"\r\n")) :]
restamped: Final = json.dumps(
{**payload, "system_fingerprint": fingerprint}, separators=(",", ":"), ensure_ascii=False
)
return _SSE_DATA_FIELD + b" " + restamped.encode() + terminator
def _awaits_fingerprint(chunk: _RelayedChunk | None) -> bool:
return (
chunk is not None
and chunk.usage is None
and chunk.system_fingerprint is None
and any(choice.finish_reason is not None for choice in chunk.choices)
)
def _is_usage_only(chunk: _RelayedChunk) -> bool:
return chunk.usage is not None and not chunk.choices
def _relayed_frame(frame: bytes, state: _RelayState) -> tuple[bytes, _RelayState]:
if state.passes_usage_through:
return frame, state
chunk: Final = _relayed_chunk(frame)
if chunk is not None and _is_usage_only(chunk):
return (
b"".join(_with_fingerprint(held, chunk.system_fingerprint) for held in _held_frames(state)),
_RelayState(),
)
if _awaits_fingerprint(chunk):
return b"", replace(state, newest_held=_Stacked(frame, state.newest_held))
return (
b"".join((*_held_frames(state), frame)),
_RelayState(passes_usage_through=chunk is not None and chunk.usage is not None),
)
def _sse_frames(complete: bytes) -> tuple[bytes, ...]:
ends: Final = tuple(match.end() for match in _SSE_FRAME_END.finditer(complete))
return tuple(complete[start:end] for start, end in zip((0, *ends[:-1]), ends))
async def without_injected_usage_frames(stream: AsyncGenerator[bytes, None]) -> AsyncIterator[bytes]:
unterminated: _Stacked | None = None # rebind-ok: stacks reads until one ends a frame, joined once
state = _RelayState() # rebind-ok: finish frames wait for the fingerprint vLLM moves onto the usage frame
try:
async for piece in stream:
if not piece:
continue
if not _ends_a_frame(unterminated, piece):
unterminated = _Stacked(piece, unterminated)
continue
complete, rest = split_complete_sse_frames(b"".join((*_oldest_first(unterminated), piece)))
unterminated = _Stacked(rest, None) if rest else None
for frame in _sse_frames(complete):
outgoing, state = _relayed_frame(frame, state)
if outgoing:
yield outgoing
if state.newest_held is not None or unterminated is not None:
yield b"".join((*_held_frames(state), *_oldest_first(unterminated)))
finally:
await stream.aclose()

34
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -79,7 +79,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul
user_api_key_auth_websocket,
user_api_key_auth_websocket_for_model,
)
from litellm.proxy.common_request_processing import open_sse_before_first_byte
from litellm.proxy.common_request_processing import model_deployments_support_stream_options, open_sse_before_first_byte
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -96,6 +96,10 @@ from litellm.proxy.common_utils.sse_keepalive import (
wrap_passthrough_sse_bytes_with_keepalive_pings,
)
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
stream_usage_relay,
without_injected_usage_frames,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
create_pass_through_route,
@ -509,6 +513,17 @@ async def vllm_proxy_route(
is_router_model: Final = is_passthrough_request_using_router_model(request_body, llm_router)
is_streaming_request: Final = is_passthrough_request_streaming(request_body)
if is_router_model and llm_router:
is_json_request: Final = is_json_content_type(request.headers.get("content-type", ""))
relay: Final = stream_usage_relay(
endpoint=endpoint,
request_body=request_body,
general_settings=_proxy_general_settings(),
supports_stream_options=lambda: model_deployments_support_stream_options(
model=request_body.get("model"),
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
),
)
result: Final = cast(
httpx.Response,
await llm_router.allm_passthrough_route(
@ -521,7 +536,7 @@ async def vllm_proxy_route(
content=None,
data=None,
files=None,
json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None),
json=(relay.body if is_json_request else None),
params=None,
headers=None,
cookies=None,
@ -530,8 +545,13 @@ async def vllm_proxy_route(
)
if is_streaming_request:
upstream_bytes: Final = _relay_response_bytes(result)
return StreamingResponse(
content=result.aiter_bytes(),
content=(
without_injected_usage_frames(upstream_bytes)
if relay.strips_usage and is_json_request
else upstream_bytes
),
status_code=result.status_code,
headers=HttpPassThroughEndpointHelpers.get_response_headers(
headers=result.headers,
@ -2058,6 +2078,14 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async
await upstream.aclose()
async def _relay_response_bytes(upstream: httpx.Response) -> AsyncGenerator[bytes, None]:
try:
async for chunk in upstream.aiter_bytes():
yield chunk
finally:
await upstream.aclose()
async def _relay_upstream_response(upstream: httpx.Response) -> Response:
return Response(
content=await upstream.aread(),

228
tests/integration/routing/test_azure_passthrough_spend_row.py Normal file
View file

@ -0,0 +1,228 @@
import json
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from typing import Final
from urllib.parse import urlsplit
import pytest
from integration._support.client import JSON_OBJECT, Gateway, held
from integration._support.database import read_rows
from integration._support.openai_wire import answering_model_discovery, responses_reply
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
_DEPLOYMENT: Final = "relay-deployment"
_API_VERSION: Final = "2024-10-21"
_INPUT_COST: Final = 0.001
_OUTPUT_COST: Final = 0.002
_USAGE: Final[dict[str, JsonValue]] = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
def _marker() -> str:
return "azure-relay-" + uuid.uuid4().hex
def _frame(payload: Mapping[str, JsonValue] | str) -> bytes:
return b"data: " + (payload if isinstance(payload, str) else json.dumps(payload)).encode() + b"\n\n"
def _chunk(identity: str, choices: Sequence[JsonValue], **extra: JsonValue) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt",
"choices": list(choices),
**extra,
}
def _chat_stream(identity: str, with_usage: bool) -> tuple[bytes, ...]:
return (
_frame(_chunk(identity, [{"index": 0, "delta": {"content": "Hello there"}, "finish_reason": None}])),
_frame(_chunk(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}])),
*((_frame(_chunk(identity, [], usage=_USAGE)),) if with_usage else ()),
_frame("[DONE]"),
)
def _chat_completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}],
"usage": _USAGE,
}
).encode()
def _embedding() -> bytes:
return json.dumps(
{
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}],
"model": "embed",
"usage": {"prompt_tokens": 5, "total_tokens": 5},
}
).encode()
def _asks_usage(body: Mapping[str, JsonValue]) -> bool:
options: Final = body.get("stream_options")
return isinstance(options, dict) and options.get("include_usage") is True
def _respond(marker: str) -> Callable[[Request], Reply]:
identity: Final = f"chatcmpl-{marker}"
def respond(request: Request) -> Reply:
body: Final = JSON_OBJECT.validate_json(request.body)
assert marker in json.dumps(body), body
path: Final = urlsplit(request.target).path
if path.endswith("/embeddings"):
return Reply(body=_embedding())
if path.endswith("/responses"):
return responses_reply(f"resp_{marker}", "gpt", "Hello", stream=body.get("stream") is True)
if body.get("stream") is True:
return Reply(content_type="text/event-stream", chunks=_chat_stream(identity, _asks_usage(body)))
return Reply(body=_chat_completion(identity))
return answering_model_discovery(respond)
@contextmanager
def _deployment(gateway: Gateway, marker: str, provider: str) -> Iterator[tuple[Wire, str]]:
with wire_server(_respond(marker)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"{provider}/{_DEPLOYMENT}",
api_base=wire.url,
api_key="sk-azure-relay",
api_version=_API_VERSION,
input_cost_per_token=_INPUT_COST,
output_cost_per_token=_OUTPUT_COST,
)
yield wire, model
def _posted(wire: Wire) -> tuple[tuple[str, dict[str, JsonValue]], ...]:
return tuple(
(urlsplit(request.target).path, JSON_OBJECT.validate_json(request.body))
for request in wire.drain()
if request.method == "POST"
)
def _one_row(model: str) -> dict[str, JsonValue]:
rows: Final = held(
lambda: read_rows(
"SELECT request_id, status, custom_llm_provider, prompt_tokens, completion_tokens, spend "
'FROM "LiteLLM_SpendLogs" WHERE model_group=%s',
(model,),
),
lambda found: len(found) >= 1,
holding=2,
seconds=70,
)
assert len(rows) == 1, rows
return rows[0]
def _assert_row(row: Mapping[str, JsonValue], provider: str, prompt: int, completion: int) -> None:
assert (row["status"], row["custom_llm_provider"]) == ("success", provider), row
assert (row["prompt_tokens"], row["completion_tokens"]) == (prompt, completion), row
assert row["spend"] == pytest.approx(prompt * _INPUT_COST + completion * _OUTPUT_COST), row
def _chat_path(model: str, route: str) -> str:
return (
f"/azure/openai/deployments/{model}/chat/completions?api-version={_API_VERSION}"
if route == "azure"
else f"/azure_ai/{model}/chat/completions"
)
def _chat_body(marker: str, **extra: JsonValue) -> dict[str, JsonValue]:
return {"messages": [{"role": "user", "content": marker}], **extra}
@pytest.mark.parametrize("provider", ["azure", "azure_ai"])
def test_router_model_chat_non_stream_writes_the_upstream_usage(gateway: Gateway, provider: str) -> None:
marker: Final = _marker()
with _deployment(gateway, marker, provider) as (wire, model):
response: Final = gateway.request("POST", _chat_path(model, provider), _chat_body(marker))
assert response.status_code == 200, response.text
assert response.content == _chat_completion(f"chatcmpl-{marker}"), response.text
((path, _),) = _posted(wire)
assert path.endswith("/chat/completions"), path
row: Final = _one_row(model)
assert row["request_id"] == f"chatcmpl-{marker}", row
_assert_row(row, provider, 5, 3)
@pytest.mark.parametrize("provider", ["azure", "azure_ai"])
def test_router_model_chat_stream_with_usage_is_relayed_verbatim_and_logged(gateway: Gateway, provider: str) -> None:
marker: Final = _marker()
with _deployment(gateway, marker, provider) as (wire, model):
sent: Final = _chat_body(marker, stream=True, stream_options={"include_usage": True})
response: Final = gateway.request("POST", _chat_path(model, provider), sent)
assert response.status_code == 200, response.text
assert response.content == b"".join(_chat_stream(f"chatcmpl-{marker}", True)), response.text
((_, posted),) = _posted(wire)
assert posted == sent, posted
_assert_row(_one_row(model), provider, 5, 3)
@pytest.mark.parametrize(
"extra",
[pytest.param({}, id="plain"), pytest.param({"prompt": 5}, id="stray_prompt_field")],
)
def test_azure_chat_stream_without_usage_is_forwarded_untouched_and_counted_from_the_messages(
gateway: Gateway, extra: Mapping[str, JsonValue]
) -> None:
marker: Final = _marker()
with _deployment(gateway, marker, "azure") as (wire, model):
sent: Final = _chat_body(marker, stream=True, **extra)
response: Final = gateway.request("POST", _chat_path(model, "azure"), sent)
assert response.status_code == 200, response.text
assert response.content == b"".join(_chat_stream(f"chatcmpl-{marker}", False)), response.text
((_, posted),) = _posted(wire)
assert posted == sent, posted
row: Final = _one_row(model)
prompt, completion = row["prompt_tokens"], row["completion_tokens"]
assert row["status"] == "success", row
assert isinstance(prompt, int) and prompt > 0, row
assert isinstance(completion, int) and completion > 0, row
assert row["spend"] == pytest.approx(prompt * _INPUT_COST + completion * _OUTPUT_COST), row
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
def test_azure_responses_body_model_writes_the_upstream_usage(gateway: Gateway, stream: bool) -> None:
marker: Final = _marker()
with _deployment(gateway, marker, "azure") as (wire, model):
sent: Final = {"model": model, "input": marker, **({"stream": True} if stream else {})}
response: Final = gateway.request("POST", f"/azure/openai/responses?api-version={_API_VERSION}", sent)
assert response.status_code == 200, response.text
assert f"resp_{marker}" in response.text, response.text
((path, _),) = _posted(wire)
assert path.endswith("/responses"), path
_assert_row(_one_row(model), "azure", 5, 3)
def test_azure_embeddings_write_the_upstream_prompt_tokens(gateway: Gateway) -> None:
marker: Final = _marker()
with _deployment(gateway, marker, "azure") as (wire, model):
response: Final = gateway.request(
"POST", f"/azure/openai/deployments/{model}/embeddings?api-version={_API_VERSION}", {"input": marker}
)
assert response.status_code == 200, response.text
assert response.content == _embedding(), response.text
((path, _),) = _posted(wire)
assert path == f"/openai/deployments/{_DEPLOYMENT}/embeddings", path
row: Final = _one_row(model)
assert (row["status"], row["custom_llm_provider"]) == ("success", "azure"), row
assert (row["prompt_tokens"], row["completion_tokens"]) == (5, 0), row
assert row["spend"] == pytest.approx(5 * _INPUT_COST), row

170
tests/integration/routing/test_vllm_passthrough_spend_row.py Normal file
View file

@ -0,0 +1,170 @@
import json
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final
import pytest
from integration._support.anthropic_sse import user_prompt
from integration._support.client import JSON_OBJECT, Gateway, eventually
from integration._support.database import read_rows
from integration._support.openai_wire import answering_model_discovery, chat_reply
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
_SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct"
_TEXT: Final = "Hello"
_INPUT_COST: Final = 0.001
_OUTPUT_COST: Final = 0.002
_PROMPT_TOKENS: Final = 5
_COMPLETION_TOKENS: Final = 3
_USAGE: Final[dict[str, JsonValue]] = {
"prompt_tokens": _PROMPT_TOKENS,
"completion_tokens": _COMPLETION_TOKENS,
"total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
}
_FINGERPRINT: Final = "fp-vllm-wire"
@dataclass(frozen=True, slots=True)
class Case:
target: str
request: Mapping[str, JsonValue]
relayed_stream_options: JsonValue | None
client_sees_usage: bool
_CASES: Final = {
"chat_non_stream": Case("/v1/chat/completions", {}, None, True),
"chat_stream": Case("/v1/chat/completions", {"stream": True}, {"include_usage": True}, False),
"chat_stream_usage_requested": Case(
"/v1/chat/completions",
{"stream": True, "stream_options": {"include_usage": True}},
{"include_usage": True},
True,
),
"completions_non_stream": Case("/v1/completions", {}, None, True),
"chat_stream_continuous_usage_only": Case(
"/v1/chat/completions",
{"stream": True, "stream_options": {"continuous_usage_stats": True}},
{"include_usage": True},
False,
),
"completions_stream": Case("/v1/completions", {"stream": True}, {"include_usage": True}, False),
}
def _data_frame(frame: Mapping[str, JsonValue]) -> bytes:
return b"data: " + json.dumps(frame).encode() + b"\n\n"
def _vllm_stream(identity: str, head: Mapping[str, JsonValue], content: JsonValue, finish: JsonValue) -> Reply:
chunk: Final = {"id": identity, "created": 1, "model": _SERVED_MODEL, **head}
usage_frame: Final = _data_frame({**chunk, "choices": [], "usage": _USAGE, "system_fingerprint": _FINGERPRINT})
return Reply(
content_type="text/event-stream",
chunks=(
_data_frame({**chunk, "choices": [content]}),
_data_frame({**chunk, "choices": [finish]}),
usage_frame + b"data: [DONE]\n\n",
),
)
def _chat_reply(identity: str, body: Mapping[str, object]) -> Reply:
if body.get("stream") is not True:
return chat_reply(identity, _SERVED_MODEL, _TEXT, stream=False)
return _vllm_stream(
identity,
{"object": "chat.completion.chunk"},
{"index": 0, "delta": {"content": _TEXT}, "finish_reason": None},
{"index": 0, "delta": {}, "finish_reason": "stop"},
)
def _completion_reply(identity: str, body: Mapping[str, object]) -> Reply:
if body.get("stream") is True:
return _vllm_stream(
identity,
{"object": "text_completion"},
{"index": 0, "text": _TEXT, "finish_reason": None},
{"index": 0, "text": "", "finish_reason": "stop"},
)
completion: Final = {
"id": identity,
"object": "text_completion",
"created": 1,
"model": _SERVED_MODEL,
"choices": [{"index": 0, "text": _TEXT, "finish_reason": "stop"}],
"usage": _USAGE,
}
return Reply(body=json.dumps(completion).encode())
def _vllm_upstream(identity: str, marker: str, case: Case) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert (request.method, request.target) == ("POST", case.target), request
body: Final = JSON_OBJECT.validate_json(request.body)
if case.target.endswith("/chat/completions"):
assert user_prompt(body) == marker, body
return _chat_reply(identity, body)
assert body.get("prompt") == marker, body
return _completion_reply(identity, body)
return answering_model_discovery(respond)
def _spend_rows(request_id: str) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(
"SELECT status, call_type, model, model_group, custom_llm_provider, prompt_tokens, completion_tokens, spend "
'FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(request_id,),
),
lambda rows: len(rows) >= 1,
seconds=70,
)
@pytest.mark.parametrize("name", list(_CASES))
def test_vllm_passthrough_to_a_router_model_writes_one_spend_row_with_the_upstream_usage(
gateway: Gateway, name: str
) -> None:
case: Final = _CASES[name]
marker: Final = "vllm-spend-" + uuid.uuid4().hex
identity: Final = f"cmpl-{marker}"
prompt: Final[Mapping[str, JsonValue]] = (
{"messages": [{"role": "user", "content": marker}]} if "chat" in case.target else {"prompt": marker}
)
with wire_server(_vllm_upstream(identity, marker, case)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"hosted_vllm/{_SERVED_MODEL}",
api_base=wire.url + "/v1",
input_cost_per_token=_INPUT_COST,
output_cost_per_token=_OUTPUT_COST,
)
response: Final = gateway.request(
"POST", f"/vllm{case.target[len('/v1') :]}", {"model": model, **prompt, **case.request}
)
assert response.status_code == 200, response.text
assert identity in response.text, response.text
assert ('"usage"' in response.text) is case.client_sees_usage, response.text
if case.request.get("stream") is True:
assert _FINGERPRINT in response.text, response.text
posted: Final = tuple(request for request in wire.drain() if request.method == "POST")
assert tuple(request.target for request in posted) == (case.target,)
assert JSON_OBJECT.validate_json(posted[0].body).get("stream_options") == case.relayed_stream_options
rows: Final = _spend_rows(identity)
assert len(rows) == 1, rows
assert rows[0] == {
"status": "success",
"call_type": "allm_passthrough_route",
"model": f"hosted_vllm/{_SERVED_MODEL}",
"model_group": model,
"custom_llm_provider": "hosted_vllm",
"prompt_tokens": _PROMPT_TOKENS,
"completion_tokens": _COMPLETION_TOKENS,
"spend": pytest.approx(_PROMPT_TOKENS * _INPUT_COST + _COMPLETION_TOKENS * _OUTPUT_COST),
}, rows

800
tests/integration/routing/test_vllm_passthrough_stream_relay.py Normal file
View file

@ -0,0 +1,800 @@
import asyncio
import itertools
import json
import re
import threading
import time
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Final
import httpx
import openai
import psutil
import pytest
from integration._support.anthropic_sse import user_prompt
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, held
from integration._support.database import read_rows
from integration._support.openai_wire import answering_model_discovery, openai_error, responses_reply
from integration._support.wire import Reply, Request, Wire, wire_server
from openai.types.chat import ChatCompletionChunk
from pydantic import JsonValue
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
_SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct"
_INPUT_COST: Final = 0.001
_OUTPUT_COST: Final = 0.002
_PROMPT_TOKENS: Final = 5
_COMPLETION_TOKENS: Final = 3
_USAGE: Final[dict[str, JsonValue]] = {
"prompt_tokens": _PROMPT_TOKENS,
"completion_tokens": _COMPLETION_TOKENS,
"total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
}
_FINGERPRINT: Final = "fp-vllm-relay"
_SSE_END: Final = re.compile(rb"\r\n\r\n|\n\n|\r\r")
_LF: Final = b"\n\n"
_CRLF: Final = b"\r\n\r\n"
_ROW_COLUMNS: Final = (
"request_id, status, call_type, model_group, custom_llm_provider, prompt_tokens, completion_tokens, spend, team_id"
)
@dataclass(frozen=True, slots=True)
class _Call:
marker: str
identity: str
def _call() -> _Call:
marker: Final = "vllm-relay-" + uuid.uuid4().hex
return _Call(marker, f"cmpl-{marker}")
def _frame(payload: Mapping[str, JsonValue] | str, end: bytes = _LF) -> bytes:
data: Final = payload if isinstance(payload, str) else json.dumps(payload)
return b"data: " + data.encode() + end
def _chunk(identity: str, choices: Sequence[JsonValue], **extra: JsonValue) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _SERVED_MODEL,
"choices": list(choices),
**extra,
}
def _delta(index: int, content: str) -> dict[str, JsonValue]:
return {"index": index, "delta": {"content": content}, "finish_reason": None}
def _finish(index: int) -> dict[str, JsonValue]:
return {"index": index, "delta": {}, "finish_reason": "stop"}
def _usage_only(identity: str) -> dict[str, JsonValue]:
return _chunk(identity, [], usage=_USAGE, system_fingerprint=_FINGERPRINT)
def _stamped(chunk: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
return {**chunk, "system_fingerprint": _FINGERPRINT}
def _sse(chunks: Sequence[bytes]) -> Reply:
return Reply(content_type="text/event-stream", chunks=tuple(chunks))
def _block_payload(block: bytes) -> bytes:
return b"\n".join(line[len(b"data:") :].strip() for line in block.splitlines() if line.startswith(b"data:"))
def _frames(raw: bytes) -> tuple[JsonValue, ...]:
payloads: Final = tuple(_block_payload(block) for block in _SSE_END.split(raw))
return tuple("[DONE]" if payload == b"[DONE]" else json.loads(payload) for payload in payloads if payload)
def _asks_usage(body: Mapping[str, JsonValue]) -> bool:
options: Final = body.get("stream_options")
return isinstance(options, dict) and options.get("include_usage") in (True, "true")
def _plain_stream(call: _Call, body: Mapping[str, JsonValue]) -> Reply:
return _sse(
(
_frame(_chunk(call.identity, [_delta(0, "Hello")])),
_frame(_chunk(call.identity, [_finish(0)])),
*((_frame(_usage_only(call.identity)),) if _asks_usage(body) else ()),
_frame("[DONE]"),
)
)
def _chat_body(model: str, call: _Call, **extra: JsonValue) -> dict[str, JsonValue]:
return {"model": model, "messages": [{"role": "user", "content": call.marker}], **extra}
def _scripted(call: _Call, reply: Callable[[Mapping[str, JsonValue]], Reply]) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST", request
body: Final = JSON_OBJECT.validate_json(request.body)
assert call.marker in json.dumps(body), body
return reply(body)
return answering_model_discovery(respond)
def _deploy(scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model=f"hosted_vllm/{_SERVED_MODEL}",
api_base=wire.url + "/v1",
input_cost_per_token=_INPUT_COST,
output_cost_per_token=_OUTPUT_COST,
)
@contextmanager
def _deployment(gateway: Gateway, respond: Callable[[Request], Reply]) -> Iterator[tuple[Wire, str]]:
with wire_server(respond) as wire, gateway.scenario() as scenario:
yield wire, _deploy(scenario, wire)
def _posted(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
return tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST")
def _read_rows(model: str) -> list[dict[str, JsonValue]]:
return read_rows(f'SELECT {_ROW_COLUMNS} FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,))
def _rows(model: str, count: int = 1) -> list[dict[str, JsonValue]]:
return held(lambda: _read_rows(model), lambda rows: len(rows) >= count, holding=2, seconds=70)
def _assert_reported_usage(row: Mapping[str, JsonValue]) -> None:
assert row["status"] == "success", row
assert row["custom_llm_provider"] == "hosted_vllm", row
assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _COMPLETION_TOKENS), row
assert row["spend"] == pytest.approx(_PROMPT_TOKENS * _INPUT_COST + _COMPLETION_TOKENS * _OUTPUT_COST), row
def _assert_counted_usage(row: Mapping[str, JsonValue]) -> None:
prompt: Final = row["prompt_tokens"]
completion: Final = row["completion_tokens"]
assert row["status"] == "success", row
assert isinstance(prompt, int) and isinstance(completion, int) and prompt > 0 and completion > 0, row
assert row["spend"] == pytest.approx(prompt * _INPUT_COST + completion * _OUTPUT_COST), row
def _one_row(model: str, identity: str | None = None) -> dict[str, JsonValue]:
rows: Final = _rows(model)
assert len(rows) == 1, rows
assert identity is None or rows[0]["request_id"] == identity, rows
return rows[0]
def _stream(gateway: Gateway, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> httpx.Response:
response: Final = gateway.request("POST", path, body, key=key)
assert response.status_code == 200, response.text
return response
def test_two_choice_stream_keeps_every_frame_in_order_and_stamps_only_the_last_finish(gateway: Gateway) -> None:
call: Final = _call()
chunks: Final = (
_chunk(call.identity, [_delta(0, "A")]),
_chunk(call.identity, [_delta(1, "B")]),
_chunk(call.identity, [_finish(0)]),
_chunk(call.identity, [_delta(1, "C")]),
_chunk(call.identity, [_finish(1)]),
)
reply: Final = _sse((*map(_frame, chunks), _frame(_usage_only(call.identity)), _frame("[DONE]")))
with _deployment(gateway, _scripted(call, lambda _: reply)) as (wire, model):
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, n=2))
assert _frames(response.content) == (*chunks[:4], _stamped(chunks[4]), "[DONE]"), response.text
(posted,) = _posted(wire)
assert posted["stream_options"] == {"include_usage": True}, posted
_assert_reported_usage(_one_row(model, call.identity))
def test_stream_with_usage_on_every_choice_chunk_is_relayed_byte_for_byte(gateway: Gateway) -> None:
call: Final = _call()
running: Final = tuple(
{"prompt_tokens": _PROMPT_TOKENS, "completion_tokens": done, "total_tokens": _PROMPT_TOKENS + done}
for done in (1, 2, 3)
)
upstream_bytes: Final = (
_frame(_chunk(call.identity, [_delta(0, "Hel")], usage=running[0])),
_frame(_chunk(call.identity, [_delta(0, "lo")], usage=running[1])),
_frame(_chunk(call.identity, [_finish(0)], usage=running[2])),
_frame(_usage_only(call.identity)),
_frame("[DONE]"),
)
with _deployment(gateway, _scripted(call, lambda _: _sse(upstream_bytes))) as (_, model):
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
assert response.content == b"".join(upstream_bytes), response.text
_assert_reported_usage(_one_row(model, call.identity))
def test_stream_split_into_single_bytes_is_reassembled_and_stripped(gateway: Gateway) -> None:
call: Final = _call()
content: Final = _chunk(call.identity, [_delta(0, "Hello")])
finish: Final = _chunk(call.identity, [_finish(0)])
whole: Final = b"".join((_frame(content), _frame(finish), _frame(_usage_only(call.identity)), _frame("[DONE]")))
reply: Final = _sse(tuple(whole[index : index + 1] for index in range(len(whole))))
with _deployment(gateway, _scripted(call, lambda _: reply)) as (_, model):
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
assert _frames(response.content) == (content, _stamped(finish), "[DONE]"), response.text
_assert_reported_usage(_one_row(model, call.identity))
def test_crlf_terminated_frames_keep_their_terminators_after_the_usage_frame_is_removed(gateway: Gateway) -> None:
call: Final = _call()
content: Final = _chunk(call.identity, [_delta(0, "Hello")])
finish: Final = _chunk(call.identity, [_finish(0)])
reply: Final = _sse(
tuple(_frame(payload, _CRLF) for payload in (content, finish, _usage_only(call.identity), "[DONE]"))
)
with _deployment(gateway, _scripted(call, lambda _: reply)) as (_, model):
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
assert _frames(response.content) == (content, _stamped(finish), "[DONE]"), response.text
assert response.content.endswith(_CRLF) and response.content.count(_CRLF) == 3, response.content
assert _LF not in response.content.replace(_CRLF, b""), response.content
_assert_reported_usage(_one_row(model, call.identity))
def test_stream_whose_upstream_sends_no_usage_frame_is_relayed_whole_and_counted_locally(gateway: Gateway) -> None:
call: Final = _call()
upstream_bytes: Final = (
_frame(_chunk(call.identity, [_delta(0, "Hello there")])),
_frame(_chunk(call.identity, [_finish(0)])),
_frame("[DONE]"),
)
with _deployment(gateway, _scripted(call, lambda _: _sse(upstream_bytes))) as (wire, model):
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
assert response.content == b"".join(upstream_bytes), response.text
(posted,) = _posted(wire)
assert posted["stream_options"] == {"include_usage": True}, posted
row: Final = _one_row(model, call.identity)
_assert_counted_usage(row)
def test_upstream_error_on_a_stream_fails_the_call_and_the_model_keeps_serving(gateway: Gateway) -> None:
call: Final = _call()
replies: Final = iter((openai_error(400), _plain_stream(call, {"stream_options": {"include_usage": True}})))
with _deployment(gateway, _scripted(call, lambda _: next(replies))) as (_, model):
eventually(lambda: _all_workers_serve(gateway, model), bool, seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS)
failed: Final = gateway.request(
"POST", "/vllm/chat/completions", _chat_body(model, call, stream=True), headers={"Connection": "close"}
)
assert failed.status_code >= 400, failed.text
assert call.identity not in failed.text, failed.text
served: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
assert _frames(served.content)[-1] == "[DONE]", served.text
_MALFORMED_STREAM_OPTIONS: Final[Mapping[str, JsonValue]] = {
"string": "include_usage",
"list": ["include_usage"],
"integer": 1,
"five_kb_string": "u" * 5000,
}
def _rejects_malformed_stream_options(call: _Call) -> Callable[[Mapping[str, JsonValue]], Reply]:
def reply(body: Mapping[str, JsonValue]) -> Reply:
if isinstance(body.get("stream_options"), dict):
return _plain_stream(call, body)
return openai_error(400)
return reply
@pytest.mark.parametrize("shape", list(_MALFORMED_STREAM_OPTIONS))
def test_malformed_stream_options_reach_the_upstream_as_the_caller_sent_them(gateway: Gateway, shape: str) -> None:
call: Final = _call()
sent: Final = _MALFORMED_STREAM_OPTIONS[shape]
with _deployment(gateway, _scripted(call, _rejects_malformed_stream_options(call))) as (wire, model):
response: Final = gateway.request(
"POST",
"/vllm/chat/completions",
_chat_body(model, call, stream=True, stream_options=sent),
headers={"Connection": "close"},
)
(posted,) = _posted(wire)
assert posted["stream_options"] == sent, posted
assert response.status_code >= 400, response.text
def test_null_stream_options_get_usage_injected_and_the_usage_frame_removed(gateway: Gateway) -> None:
call: Final = _call()
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
response: Final = _stream(
gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, stream_options=None)
)
(posted,) = _posted(wire)
assert posted["stream_options"] == {"include_usage": True}, posted
assert _frames(response.content) == (
_chunk(call.identity, [_delta(0, "Hello")]),
_stamped(_chunk(call.identity, [_finish(0)])),
"[DONE]",
), response.text
_assert_reported_usage(_one_row(model, call.identity))
@pytest.mark.parametrize(
"sent",
[pytest.param({}, id="empty"), pytest.param({"include_usage": False}, id="include_usage_false")],
)
def test_stream_options_without_include_usage_true_get_usage_injected_and_the_usage_frame_removed(
gateway: Gateway, sent: Mapping[str, JsonValue]
) -> None:
call: Final = _call()
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
response: Final = _stream(
gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, stream_options=dict(sent))
)
(posted,) = _posted(wire)
assert posted["stream_options"] == {"include_usage": True}, posted
assert _frames(response.content) == (
_chunk(call.identity, [_delta(0, "Hello")]),
_stamped(_chunk(call.identity, [_finish(0)])),
"[DONE]",
), response.text
_assert_reported_usage(_one_row(model, call.identity))
def test_null_include_usage_reaches_the_upstream_as_sent_and_is_counted_locally(gateway: Gateway) -> None:
call: Final = _call()
sent: Final[dict[str, JsonValue]] = {"include_usage": None}
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
response: Final = _stream(
gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, stream_options=sent)
)
(posted,) = _posted(wire)
assert posted["stream_options"] == sent, posted
assert response.content == b"".join(_plain_stream(call, posted).chunks or ()), response.text
_assert_counted_usage(_one_row(model, call.identity))
def test_duplicated_stream_options_key_takes_the_last_value(gateway: Gateway) -> None:
call: Final = _call()
raw: Final = (
'{"model": "%s", "messages": [{"role": "user", "content": "%s"}], "stream": true, '
'"stream_options": {"include_usage": true}, "stream_options": {"continuous_usage_stats": true}}'
)
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
response: Final = gateway.client.post(
"/vllm/chat/completions",
content=(raw % (model, call.marker)).encode(),
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
)
assert response.status_code == 200, response.text
(posted,) = _posted(wire)
assert posted["stream_options"] == {"include_usage": True}, posted
assert '"usage"' not in response.text, response.text
_assert_reported_usage(_one_row(model, call.identity))
def test_include_usage_sent_as_the_string_true_still_reaches_the_caller(gateway: Gateway) -> None:
call: Final = _call()
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (_, model):
response: Final = _stream(
gateway,
"/vllm/chat/completions",
_chat_body(model, call, stream=True, stream_options={"include_usage": "true"}),
)
assert _usage_only(call.identity) in _frames(response.content), response.text
_assert_reported_usage(_one_row(model, call.identity))
def test_stream_flag_sent_as_a_string_streams_verbatim_and_is_counted_locally(gateway: Gateway) -> None:
call: Final = _call()
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream="true"))
(posted,) = _posted(wire)
assert "stream_options" not in posted, posted
assert response.content == b"".join(_plain_stream(call, posted).chunks or ()), response.text
_assert_counted_usage(_one_row(model, call.identity))
def _upstream_socket_states(wire: Wire) -> frozenset[str]:
port: Final = httpx.URL(wire.url).port
return frozenset(
connection.status
for connection in psutil.Process().net_connections(kind="tcp")
if connection.laddr.port == port
)
def test_caller_disconnect_mid_stream_closes_the_upstream_and_logs_what_was_relayed(gateway: Gateway) -> None:
call: Final = _call()
release: Final = threading.Event()
reply: Final = Reply(
content_type="text/event-stream",
chunks=(
_frame(_chunk(call.identity, [_delta(0, "Hello there")])),
_frame(_chunk(call.identity, [_finish(0)])),
_frame(_usage_only(call.identity)) + _frame("[DONE]"),
),
gate_after_first=release,
)
with _deployment(gateway, _scripted(call, lambda _: reply)) as (wire, model):
try:
with gateway.client.stream(
"POST",
"/vllm/chat/completions",
json=_chat_body(model, call, stream=True),
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200
assert call.identity.encode() in next(response.iter_raw())
eventually(
lambda: _upstream_socket_states(wire),
lambda states: psutil.CONN_CLOSE_WAIT in states and psutil.CONN_ESTABLISHED not in states,
seconds=30,
)
finally:
release.set()
_assert_counted_usage(_one_row(model, call.identity))
def _embedding_reply(_: Mapping[str, JsonValue]) -> Reply:
body: Final = {
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5, 0.75]}],
"model": _SERVED_MODEL,
"usage": {"prompt_tokens": _PROMPT_TOKENS, "total_tokens": _PROMPT_TOKENS},
}
return Reply(body=json.dumps(body).encode())
def test_embeddings_write_one_row_with_the_upstream_prompt_tokens(gateway: Gateway) -> None:
call: Final = _call()
with _deployment(gateway, _scripted(call, _embedding_reply)) as (wire, model):
response: Final = gateway.request("POST", "/vllm/embeddings", {"model": model, "input": call.marker})
assert response.status_code == 200, response.text
assert JSON_OBJECT.validate_json(response.content)["data"] == [
{"object": "embedding", "index": 0, "embedding": [0.25, 0.5, 0.75]}
], response.text
assert tuple(request.target for request in wire.drain() if request.method == "POST") == ("/v1/embeddings",)
row: Final = _one_row(model)
assert row["status"] == "success", row
assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, 0), row
assert row["spend"] == pytest.approx(_PROMPT_TOKENS * _INPUT_COST), row
@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
def test_responses_write_one_row_and_relay_the_body_untouched(gateway: Gateway, stream: bool) -> None:
call: Final = _call()
reply: Final = responses_reply(f"resp_{call.marker}", _SERVED_MODEL, "Hello", stream=stream)
with _deployment(gateway, _scripted(call, lambda _: reply)) as (wire, model):
sent: Final = {"model": model, "input": call.marker, **({"stream": True} if stream else {})}
response: Final = gateway.request("POST", "/vllm/responses", sent)
assert response.status_code == 200, response.text
assert response.content == (b"".join(reply.chunks) if reply.chunks else reply.body), response.text
(posted,) = _posted(wire)
assert {name: value for name, value in posted.items() if name != "model"} == {
name: value for name, value in sent.items() if name != "model"
}, posted
_assert_reported_usage(_one_row(model))
def _pinned(gateway: Gateway) -> httpx.Client:
return httpx.Client(
base_url=gateway.client.base_url,
timeout=30,
trust_env=False,
limits=httpx.Limits(max_connections=1, max_keepalive_connections=1, keepalive_expiry=300),
)
def _chunk_text(chunk: ChatCompletionChunk) -> str:
return "".join(choice.delta.content or "" for choice in chunk.choices)
def _listed(response: httpx.Response, model: str) -> bool:
return response.status_code == 200 and model in response.text
def test_openai_sdk_stream_assembles_the_answer_without_a_usage_chunk(gateway: Gateway) -> None:
call: Final = _call()
with (
_deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (_, model),
_pinned(gateway) as http,
):
auth: Final = {"Authorization": f"Bearer {gateway.key}"}
eventually(
lambda: http.get("/v1/models", headers=auth),
lambda listed: _listed(listed, model),
seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
)
client: Final = openai.OpenAI(
base_url=f"{str(gateway.client.base_url).rstrip('/')}/vllm",
api_key=gateway.key,
http_client=http,
max_retries=0,
)
chunks: Final = tuple(
client.chat.completions.create(
model=model, messages=[{"role": "user", "content": call.marker}], stream=True
)
)
assert "".join(map(_chunk_text, chunks)) == "Hello", chunks
assert [chunk.usage for chunk in chunks if chunk.usage is not None] == [], chunks
assert {chunk.id for chunk in chunks} == {call.identity}, chunks
_assert_reported_usage(_one_row(model, call.identity))
async def _async_listed(http: httpx.AsyncClient, key: str, model: str) -> httpx.Response:
return await http.get("/v1/models", headers={"Authorization": f"Bearer {key}"})
async def test_async_openai_sdk_stream_assembles_the_answer_without_a_usage_chunk(gateway: Gateway) -> None:
call: Final = _call()
loop: Final = asyncio.get_running_loop()
with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (_, model):
async with httpx.AsyncClient(
base_url=gateway.client.base_url,
timeout=30,
trust_env=False,
limits=httpx.Limits(max_connections=1, max_keepalive_connections=1, keepalive_expiry=300),
) as http:
await asyncio.to_thread(
eventually,
lambda: asyncio.run_coroutine_threadsafe(_async_listed(http, gateway.key, model), loop).result(),
lambda listed: _listed(listed, model),
2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
)
client: Final = openai.AsyncOpenAI(
base_url=f"{str(gateway.client.base_url).rstrip('/')}/vllm",
api_key=gateway.key,
http_client=http,
max_retries=0,
)
stream: Final = await client.chat.completions.create(
model=model, messages=[{"role": "user", "content": call.marker}], stream=True
)
chunks: Final = tuple([chunk async for chunk in stream])
assert "".join(map(_chunk_text, chunks)) == "Hello", chunks
assert [chunk.usage for chunk in chunks if chunk.usage is not None] == [], chunks
assert {chunk.id for chunk in chunks} == {call.identity}, chunks
_assert_reported_usage(_one_row(model, call.identity))
def test_identical_requests_repeated_each_reach_the_upstream_and_each_write_one_row(gateway: Gateway) -> None:
call: Final = _call()
served: Final = itertools.count()
def reply(body: Mapping[str, JsonValue]) -> Reply:
repeat: Final = _Call(call.marker, f"{call.identity}-{next(served)}")
if body.get("stream") is True:
return _plain_stream(repeat, body)
completion: Final = {
"id": repeat.identity,
"object": "chat.completion",
"created": 1,
"model": _SERVED_MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}],
"usage": _USAGE,
}
return Reply(body=json.dumps(completion).encode())
with _deployment(gateway, _scripted(call, reply)) as (wire, model):
bodies: Final = (_chat_body(model, call), _chat_body(model, call), *(_chat_body(model, call, stream=True),) * 2)
texts: Final = tuple(_stream(gateway, "/vllm/chat/completions", body).text for body in bodies)
identities: Final = tuple(f"{call.identity}-{index}" for index in range(4))
assert all(identity in text for identity, text in zip(identities, texts)), texts
assert len(_posted(wire)) == 4
rows: Final = _rows(model, 4)
assert sorted(str(row["request_id"]) for row in rows) == sorted(identities), rows
for row in rows:
_assert_reported_usage(row)
def test_non_json_content_type_never_reaches_the_upstream(gateway: Gateway) -> None:
call: Final = _call()
with (
wire_server(answering_model_discovery(lambda _: _plain_stream(call, {}))) as wire,
gateway.scenario() as scenario,
):
model: Final = _deploy(scenario, wire)
response: Final = gateway.client.post(
"/vllm/chat/completions",
content=json.dumps(_chat_body(model, call, stream=True)).encode(),
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "text/plain", "Connection": "close"},
)
posted: Final = tuple(request.body for request in wire.drain() if request.method == "POST")
assert response.status_code >= 400, response.text
assert posted == (), posted
def test_team_key_stream_writes_the_row_under_the_team(gateway: Gateway) -> None:
call: Final = _call()
with wire_server(_scripted(call, lambda body: _plain_stream(call, body))) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, wire)
team: Final = scenario.team()
key: Final = scenario.key(team_id=team)
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True), key=key)
assert '"usage"' not in response.text, response.text
row: Final = _one_row(model, call.identity)
assert row["team_id"] == team, row
_assert_reported_usage(row)
_LARGE_FRAME_BYTES: Final = 8 * 1024 * 1024
_PIECE_BYTES: Final = 8 * 1024
_HEALTH_CEILING_SECONDS: Final = 1.0
_RELAY_RATIO_CEILING: Final = 4
_RELAY_SLACK_SECONDS: Final = 1.0
@dataclass(frozen=True, slots=True)
class _Relay:
response: httpx.Response
seconds: float
health: tuple[float, ...]
def _timed_health(client: httpx.Client) -> float:
started: Final = time.perf_counter()
response: Final = client.get("/health/liveliness", headers={"Connection": "close"})
assert response.status_code == 200, response.text
return time.perf_counter() - started
def _health_latencies(gateway: Gateway, done: threading.Event) -> tuple[float, ...]:
with httpx.Client(base_url=gateway.client.base_url, timeout=60, trust_env=False) as fresh:
return tuple(_timed_health(fresh) for _ in iter(done.is_set, True))
def _finish_and_done(call: _Call) -> bytes:
return _frame(_chunk(call.identity, [_finish(0)])) + _frame("[DONE]")
def _content_then_usage(call: _Call, content: Sequence[bytes]) -> Callable[[Mapping[str, JsonValue]], Reply]:
usage: Final = _frame(_chunk(call.identity, [], usage=_USAGE))
finish: Final = _frame(_chunk(call.identity, [_finish(0)]))
def reply(body: Mapping[str, JsonValue]) -> Reply:
return _sse((*content, finish, *((usage,) if _asks_usage(body) else ()), _frame("[DONE]")))
return reply
def _timed_relay(gateway: Gateway, call: _Call, content: Sequence[bytes]) -> _Relay:
with _deployment(gateway, _scripted(call, _content_then_usage(call, content))) as (_, model):
eventually(lambda: _all_workers_serve(gateway, model), bool, seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS)
done: Final = threading.Event()
with ThreadPoolExecutor(max_workers=1) as pool:
polled: Final = pool.submit(_health_latencies, gateway, done)
started: Final = time.perf_counter()
try:
response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
finally:
done.set()
seconds: Final = time.perf_counter() - started
return _Relay(response, seconds, polled.result())
def _assert_large_frame_relays_like_small_frames(
gateway: Gateway, record_property: Callable[[str, object], None], large_call: _Call, large: bytes
) -> None:
small_call: Final = _call()
small: Final = tuple(
_frame(_chunk(small_call.identity, [_delta(0, "x" * _PIECE_BYTES)]))
for _ in range(_LARGE_FRAME_BYTES // _PIECE_BYTES)
)
pieces: Final = tuple(large[start : start + _PIECE_BYTES] for start in range(0, len(large), _PIECE_BYTES))
small_relay: Final = _timed_relay(gateway, small_call, small)
large_relay: Final = _timed_relay(gateway, large_call, pieces)
record_property("small_frames_seconds", round(small_relay.seconds, 3))
record_property("large_frame_seconds", round(large_relay.seconds, 3))
record_property("health_calls", len(large_relay.health))
record_property("health_max_seconds", round(max(large_relay.health), 3))
assert small_relay.response.content == b"".join(small) + _finish_and_done(small_call), len(
small_relay.response.content
)
assert large_relay.response.content == large + _finish_and_done(large_call), len(large_relay.response.content)
assert max(large_relay.health) < _HEALTH_CEILING_SECONDS, sorted(large_relay.health)[-5:]
assert large_relay.seconds < _RELAY_RATIO_CEILING * small_relay.seconds + _RELAY_SLACK_SECONDS, (
large_relay.seconds,
small_relay.seconds,
)
def test_frame_terminated_only_at_its_end_relays_as_fast_as_the_same_bytes_in_small_frames(
gateway: Gateway, record_property: Callable[[str, object], None]
) -> None:
large_call: Final = _call()
large: Final = _frame(_chunk(large_call.identity, [_delta(0, "x" * _LARGE_FRAME_BYTES)]))
_assert_large_frame_relays_like_small_frames(gateway, record_property, large_call, large)
def test_frame_of_usage_and_finish_reason_near_misses_relays_as_fast_as_the_same_bytes_in_small_frames(
gateway: Gateway, record_property: Callable[[str, object], None]
) -> None:
near_miss: Final[dict[str, JsonValue]] = {"usage": 0, "finish_reason": None}
count: Final = _LARGE_FRAME_BYTES // len(json.dumps(near_miss) + ", ")
large_call: Final = _call()
large: Final = _frame(_chunk(large_call.identity, [_delta(0, "")], pad=[near_miss] * count))
_assert_large_frame_relays_like_small_frames(gateway, record_property, large_call, large)
@dataclass(frozen=True, slots=True)
class _Outcome:
identity: str
status: int
text: str
def _burst_reply(dropped: Callable[[], bool]) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
body: Final = JSON_OBJECT.validate_json(request.body)
marker: Final = user_prompt(body) if "messages" in body else str(body.get("prompt"))
call: Final = _Call(marker, f"cmpl-{marker}")
if dropped():
return Reply(drop_connection=True)
if body.get("stream") is not True:
completion: Final = {
"id": call.identity,
"object": "chat.completion",
"created": 1,
"model": _SERVED_MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "Hi"}, "finish_reason": "stop"}],
"usage": _USAGE,
}
return Reply(body=json.dumps(completion).encode())
stream: Final = _plain_stream(call, body)
return Reply(content_type=stream.content_type, chunks=stream.chunks, pause_between_chunks=0.05)
return answering_model_discovery(respond)
async def _send(client: httpx.AsyncClient, key: str, body: Mapping[str, JsonValue], identity: str) -> _Outcome:
try:
response: Final = await client.post(
"/vllm/chat/completions", json=body, headers={"Authorization": f"Bearer {key}"}
)
except httpx.TransportError as error:
return _Outcome(identity, 0, repr(error))
return _Outcome(identity, response.status_code, response.text)
def _all_workers_serve(gateway: Gateway, model: str) -> bool:
auth: Final = {"Authorization": f"Bearer {gateway.key}"}
with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as fresh:
return all(_listed(fresh.get("/v1/models", headers={**auth, "Connection": "close"}), model) for _ in range(16))
async def test_burst_with_dropped_upstream_connections_lands_every_completed_call_once(gateway: Gateway) -> None:
arrivals: Final = itertools.count(1)
with wire_server(_burst_reply(lambda: next(arrivals) % 4 == 0)) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, wire)
await asyncio.to_thread(
eventually, lambda: _all_workers_serve(gateway, model), bool, 2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
)
calls: Final = tuple(_call() for _ in range(24))
bodies: Final = tuple(
_chat_body(model, call, **({"stream": True} if index % 3 else {})) for index, call in enumerate(calls)
)
async with httpx.AsyncClient(base_url=gateway.client.base_url, timeout=90, trust_env=False) as client:
outcomes: Final = await asyncio.gather(
*(_send(client, gateway.key, body, call.identity) for body, call in zip(bodies, calls))
)
completed: Final = frozenset(
outcome.identity for outcome in outcomes if outcome.status == 200 and outcome.identity in outcome.text
)
failed: Final = tuple(outcome for outcome in outcomes if outcome.identity not in completed)
assert len(completed) >= 12, outcomes
assert all(outcome.status != 200 for outcome in failed), failed
rows: Final = _rows(model, len(completed))
assert sorted(str(row["request_id"]) for row in rows) == sorted(completed), rows
health: Final = gateway.request("GET", "/health/liveliness")
assert health.status_code == 200, health.text

283
tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py Normal file
View file

@ -0,0 +1,283 @@
import asyncio
import json
import threading
import uuid
from collections.abc import Callable, Generator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final
from urllib.parse import urlsplit
import httpx
import psutil
import pytest
import yaml
from integration._support.anthropic_sse import user_prompt
from integration._support.client import JSON_OBJECT, eventually, gateway_from_environment, held, object_value
from integration._support.database import read_rows
from integration._support.openai_wire import answering_model_discovery
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
_SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct"
_INPUT_COST: Final = 0.001
_OUTPUT_COST: Final = 0.002
_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
)
_STARTED_WORKER: Final = "Started server process ["
_STARTUP_COMPLETE: Final = "Application startup complete."
_BURST: Final = 20
@dataclass(frozen=True, slots=True)
class _Rig:
wire: Wire
proxy: OwnedProxy
model: str
@property
def base_url(self) -> str:
return str(self.proxy.gateway.client.base_url).rstrip("/")
@property
def key(self) -> str:
return self.proxy.gateway.key
@dataclass(frozen=True, slots=True)
class _Outcome:
identity: str
status: int
text: str
def _identity(body: Mapping[str, JsonValue]) -> str:
return f"cmpl-{user_prompt(body)}"
def _frame(payload: Mapping[str, JsonValue] | str) -> bytes:
return b"data: " + (payload if isinstance(payload, str) else json.dumps(payload)).encode() + b"\n\n"
def _chunk(identity: str, choices: Sequence[JsonValue], **extra: JsonValue) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _SERVED_MODEL,
"choices": list(choices),
**extra,
}
def _usage_frame(identity: str) -> bytes:
return _frame(_chunk(identity, [], usage=dict(_USAGE)))
def _asks_usage(body: Mapping[str, JsonValue]) -> bool:
options: Final = body.get("stream_options")
return isinstance(options, dict) and options.get("include_usage") is True
def _stream_bytes(body: Mapping[str, JsonValue]) -> tuple[bytes, ...]:
identity: Final = _identity(body)
return (
_frame(_chunk(identity, [{"index": 0, "delta": {"content": "Hello there"}, "finish_reason": None}])),
_frame(_chunk(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}])),
*((_usage_frame(identity),) if _asks_usage(body) else ()),
_frame("[DONE]"),
)
def _streaming(request: Request) -> Reply:
return Reply(content_type="text/event-stream", chunks=_stream_bytes(JSON_OBJECT.validate_json(request.body)))
def _config(directory: Path, wire: Wire, model: str, general_settings: Mapping[str, JsonValue]) -> Path:
base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
deployment: Final = {
"model_name": model,
"litellm_params": {
"model": f"hosted_vllm/{_SERVED_MODEL}",
"api_base": f"{wire.url}/v1",
"input_cost_per_token": _INPUT_COST,
"output_cost_per_token": _OUTPUT_COST,
},
}
config: Final = {
**base,
"model_list": [*(base.get("model_list") or []), deployment],
"general_settings": {**object_value(base["general_settings"]), **general_settings},
}
path: Final = directory / f"vllm-usage-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@contextmanager
def _vllm_proxy(
directory: Path, respond: Callable[[Request], Reply], general_settings: Mapping[str, JsonValue]
) -> Generator[_Rig, None, None]:
model: Final = f"vllm-usage-{uuid.uuid4().hex}"
with gateway_from_environment() as gateway, wire_server(answering_model_discovery(respond)) as wire:
config: Final = _config(directory, wire, model, general_settings)
with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
yield _Rig(wire, owned, model)
def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]:
return {"model": model, "messages": [{"role": "user", "content": f"vllm-usage-{uuid.uuid4().hex}"}], **extra}
def _rows(model: str, count: int) -> list[dict[str, JsonValue]]:
return held(
lambda: read_rows(
'SELECT request_id, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
"WHERE model_group=%s",
(model,),
),
lambda rows: len(rows) >= count,
holding=2,
seconds=70,
)
def _posted(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
return tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST")
def _spend(prompt: JsonValue, completion: JsonValue) -> float:
assert isinstance(prompt, int) and isinstance(completion, int), (prompt, completion)
return prompt * _INPUT_COST + completion * _OUTPUT_COST
@pytest.mark.parametrize("setting", [True, None], ids=["true", "null"])
def test_always_include_stream_usage_true_or_null_injects_usage_and_relays_the_stream_the_caller_asked_for(
tmp_path: Path, setting: bool | None
) -> None:
with _vllm_proxy(tmp_path, _streaming, {"always_include_stream_usage": setting}) as rig:
body: Final = _chat_body(rig.model, stream=True)
response: Final = rig.proxy.gateway.request("POST", "/vllm/chat/completions", body)
assert response.status_code == 200, response.text
(posted,) = _posted(rig.wire)
assert posted["stream_options"] == {"include_usage": True}, posted
assert response.content == b"".join(_stream_bytes(body)), response.text
(row,) = _rows(rig.model, 1)
assert row["request_id"] == _identity(body), row
assert (row["status"], row["prompt_tokens"], row["completion_tokens"]) == ("success", 5, 3), row
assert row["spend"] == pytest.approx(_spend(5, 3)), row
def test_always_include_stream_usage_false_forwards_the_body_untouched_and_counts_locally(tmp_path: Path) -> None:
with _vllm_proxy(tmp_path, _streaming, {"always_include_stream_usage": False}) as rig:
body: Final = _chat_body(rig.model, stream=True)
response: Final = rig.proxy.gateway.request("POST", "/vllm/chat/completions", body)
assert response.status_code == 200, response.text
(posted,) = _posted(rig.wire)
assert "stream_options" not in posted, posted
assert response.content == b"".join(_stream_bytes(body)), response.text
(row,) = _rows(rig.model, 1)
assert row["request_id"] == _identity(body), row
prompt, completion = row["prompt_tokens"], row["completion_tokens"]
assert row["status"] == "success" and isinstance(prompt, int) and prompt > 0, row
assert isinstance(completion, int) and completion > 0, row
assert row["spend"] == pytest.approx(_spend(prompt, completion)), row
def _holding(release: threading.Event, arrived: SimpleQueue[str]) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method == "POST" and urlsplit(request.target).path.endswith("/chat/completions"):
arrived.put(request.target)
assert release.wait(timeout=90), "The burst was never released"
return _streaming(request)
return respond
@contextmanager
def _releasing(release: threading.Event) -> Generator[None, None, None]:
try:
yield
finally:
release.set()
def _worker_pids(log: str) -> tuple[int, ...]:
return tuple(
int(line.split(_STARTED_WORKER, 1)[1].split("]", 1)[0]) for line in log.splitlines() if _STARTED_WORKER in line
)
def _open_upstream_connections(pid: int, upstream: str) -> int:
port: Final = urlsplit(upstream).port
return sum(
1
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
)
async def _send(client: httpx.AsyncClient, key: str, body: Mapping[str, JsonValue]) -> _Outcome:
try:
response: Final = await client.post(
"/vllm/chat/completions", json=body, headers={"Authorization": f"Bearer {key}", "Connection": "close"}
)
except httpx.TransportError as error:
return _Outcome(_identity(body), 0, repr(error))
return _Outcome(_identity(body), response.status_code, response.text)
async def _burst(rig: _Rig, bodies: Sequence[Mapping[str, JsonValue]]) -> tuple[_Outcome, ...]:
async with httpx.AsyncClient(base_url=rig.base_url, timeout=120, trust_env=False) as client:
return tuple(await asyncio.gather(*(_send(client, rig.key, body) for body in bodies)))
def _served(outcomes: Sequence[_Outcome]) -> frozenset[str]:
return frozenset(
outcome.identity for outcome in outcomes if outcome.status == 200 and outcome.text.endswith("data: [DONE]\n\n")
)
async def test_worker_sigkill_mid_burst_lands_every_served_call_exactly_once(tmp_path: Path) -> None:
release: Final = threading.Event()
arrived: Final[SimpleQueue[str]] = SimpleQueue()
with _vllm_proxy(tmp_path, _holding(release, arrived), {}) as rig, _releasing(release):
workers: Final = eventually(
lambda: _worker_pids(rig.proxy.log.read_text()), lambda pids: len(pids) == 2, seconds=30
)
burst: Final = asyncio.create_task(
_burst(rig, tuple(_chat_body(rig.model, stream=True) for _ in range(_BURST)))
)
await asyncio.to_thread(eventually, arrived.qsize, lambda size: size == _BURST, 60)
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, rig.wire.url) for pid in workers})
assert sum(held_by.values()) == _BURST, held_by
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
with rig.proxy.replacing_worker(victim_pid) as kill:
kill.freeze()
await asyncio.to_thread(kill)
release.set()
outcomes: Final = await burst
served: Final = _served(outcomes)
assert held_by[survivor_pid] >= _BURST // 2, held_by
assert len(served) == held_by[survivor_pid], (held_by, outcomes)
assert all(outcome.status != 200 for outcome in outcomes if outcome.identity not in served), outcomes
assert all('"usage"' not in outcome.text for outcome in outcomes if outcome.identity in served), outcomes
(follow_up,) = await _burst(rig, (_chat_body(rig.model, stream=True),))
assert _served((follow_up,)) == {follow_up.identity}, follow_up
rows: Final = _rows(rig.model, len(served) + 1)
assert sorted(str(row["request_id"]) for row in rows) == sorted((*served, follow_up.identity)), rows
assert all(row["status"] == "success" for row in rows), rows
replaced: Final = eventually(
lambda: rig.proxy.log.read_text(),
lambda log: len(_worker_pids(log)) == 3 and log.count(_STARTUP_COMPLETE) == 3,
seconds=2 * graceful_stop_seconds(),
)
assert victim_pid not in _worker_pids(replaced)[2:], replaced

5
tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py
View file

@ -230,10 +230,11 @@ def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response():
assert response.usage.completion_tokens == 8
def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request():
@pytest.mark.parametrize("stray_fields", [{}, {"prompt": 5}])
def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(stray_fields):
messages = [{"role": "user", "content": "Say hi in three words"}]
logging_obj = MagicMock()
logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}}
logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True, **stray_fields}}
response = AzurePassthroughConfig().handle_logging_collected_chunks(
all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk],

0
tests/unit/llms/vllm/__init__.py Normal file
View file

0
tests/unit/llms/vllm/passthrough/__init__.py Normal file
View file

241
tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py Normal file
View file

@ -0,0 +1,241 @@
import json
from datetime import datetime
import httpx
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.vllm.passthrough.transformation import VLLMPassthroughConfig
MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
MESSAGES = [{"role": "user", "content": "Say hi in three words"}]
PROMPT = "The capital of France is"
USAGE = {"prompt_tokens": 34, "completion_tokens": 4, "total_tokens": 38}
def _relay_logging_obj(endpoint: str, request_data: dict, callbacks: list[CustomLogger] | None = None) -> Logging:
logging_obj = Logging(
model=MODEL,
messages=[],
stream="stream" in request_data,
call_type="allm_passthrough_route",
start_time=datetime(2025, 1, 1),
litellm_call_id="call-1",
function_id="fn-1",
dynamic_async_success_callbacks=callbacks,
)
logging_obj.update_environment_variables(
model=MODEL,
litellm_params={"api_base": "http://vllm.local:8000/v1"},
optional_params={},
custom_llm_provider="hosted_vllm",
endpoint=endpoint,
request_data=request_data,
)
return logging_obj
def _relayed_json(path: str, body: dict) -> httpx.Response:
return httpx.Response(
status_code=200,
headers={"content-type": "application/json"},
content=json.dumps(body).encode(),
request=httpx.Request("POST", f"http://vllm.local:8000/{path}"),
)
def _chat_completion(identity: str) -> httpx.Response:
body = {
"id": identity,
"object": "chat.completion",
"created": 1,
"model": MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello there!"}, "finish_reason": "stop"}],
"usage": USAGE,
}
return _relayed_json("v1/chat/completions", body)
def _chat_stream(identity: str, usage: dict | None) -> bytes:
head = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": MODEL}
frames = [
{**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello"}, "finish_reason": None}]},
{**head, "choices": [{"index": 0, "delta": {"content": " there!"}, "finish_reason": "stop"}]},
*([{**head, "choices": [], "usage": usage}] if usage else []),
]
return b"".join(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames) + b"data: [DONE]\n\n"
def _text_stream(identity: str, deltas: list[tuple[int, str]], usage: dict | None) -> bytes:
head = {"id": identity, "object": "text_completion", "created": 1, "model": MODEL}
frames = [
*({**head, "choices": [{"index": index, "text": text, "finish_reason": None}]} for index, text in deltas),
*([{**head, "choices": [], "usage": usage}] if usage else []),
]
return b"".join(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames) + b"data: [DONE]\n\n"
async def _log_stream(endpoint: str, sse: bytes, request_data: dict | None = None) -> Logging:
logging_obj = _relay_logging_obj(endpoint, request_data or {"model": MODEL, "messages": MESSAGES, "stream": True})
collector = VLLMPassthroughConfig().create_stream_collector(
model=MODEL, custom_llm_provider="hosted_vllm", endpoint=endpoint
)
collector.add(sse[:57])
collector.add(sse[57:])
await logging_obj.async_flush_passthrough_collected_chunks(collector)
return logging_obj
@pytest.mark.asyncio
async def test_a_relayed_chat_completion_is_logged_with_the_usage_vllm_reported():
logging_obj = _relay_logging_obj("chat/completions", {"model": MODEL, "messages": MESSAGES})
await logging_obj.async_success_handler(result=_chat_completion("chatcmpl-vllm-1"))
payload = logging_obj.model_call_details["standard_logging_object"]
assert payload["id"] == "chatcmpl-vllm-1"
assert payload["custom_llm_provider"] == "hosted_vllm"
assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (34, 4, 38)
@pytest.mark.asyncio
async def test_a_relayed_chat_stream_is_logged_with_the_usage_chunk_vllm_sent():
logging_obj = await _log_stream("v1/chat/completions", _chat_stream("chatcmpl-vllm-2", USAGE))
payload = logging_obj.model_call_details["standard_logging_object"]
assert payload["id"] == "chatcmpl-vllm-2"
assert payload["response"]["choices"][0]["message"]["content"] == "Hello there!"
assert (payload["prompt_tokens"], payload["completion_tokens"]) == (34, 4)
@pytest.mark.asyncio
async def test_a_relayed_chat_stream_without_a_usage_chunk_counts_tokens_from_the_relayed_request():
logging_obj = await _log_stream("chat/completions", _chat_stream("chatcmpl-vllm-3", None))
payload = logging_obj.model_call_details["standard_logging_object"]
assert payload["id"] == "chatcmpl-vllm-3"
assert payload["prompt_tokens"] == litellm.token_counter(model=MODEL, messages=MESSAGES)
assert payload["completion_tokens"] > 0
@pytest.mark.asyncio
async def test_a_relayed_text_completion_is_logged_with_the_usage_vllm_reported():
logging_obj = _relay_logging_obj("v1/completions", {"model": MODEL, "prompt": PROMPT})
body = {
"id": "cmpl-vllm-4",
"object": "text_completion",
"created": 1,
"model": MODEL,
"choices": [{"index": 0, "text": " Paris.", "finish_reason": "stop", "stop_reason": None}],
"usage": USAGE,
}
result = VLLMPassthroughConfig().logging_non_streaming_response(
model=MODEL,
custom_llm_provider="hosted_vllm",
httpx_response=_relayed_json("v1/completions", body),
request_data={"model": MODEL, "prompt": PROMPT},
logging_obj=logging_obj,
endpoint="v1/completions",
)
await logging_obj.async_success_handler(result=result)
payload = logging_obj.model_call_details["standard_logging_object"]
assert payload["id"] == "cmpl-vllm-4"
assert payload["response"]["choices"][0]["text"] == " Paris."
assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (34, 4, 38)
class _ResponseRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.responses: list[object] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.responses.append(response_obj)
@pytest.mark.asyncio
async def test_a_success_callback_gets_the_text_completion_vllm_relayed_as_a_parsed_response():
recorder = _ResponseRecorder()
logging_obj = _relay_logging_obj("v1/completions", {"model": MODEL, "prompt": PROMPT}, callbacks=[recorder])
body = {
"id": "cmpl-vllm-8",
"object": "text_completion",
"created": 1,
"model": MODEL,
"choices": [{"index": 0, "text": " Paris.", "finish_reason": "stop"}],
"usage": USAGE,
}
await logging_obj.async_success_handler(result=_relayed_json("v1/completions", body))
[response] = recorder.responses
assert (response.id, response.choices[0].text, response.usage.prompt_tokens) == ("cmpl-vllm-8", " Paris.", 34)
@pytest.mark.parametrize(
"path, request_data, body",
[
(
"v1/completions",
{"model": MODEL, "prompt": PROMPT},
{
"id": "cmpl-vllm-7",
"object": "text_completion",
"created": 1,
"model": MODEL,
"choices": [{"index": 0, "text": " Paris.", "finish_reason": "stop"}],
"usage": USAGE,
},
),
(
"v1/embeddings",
{"model": MODEL, "input": [PROMPT]},
{
"object": "list",
"model": MODEL,
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"usage": {"prompt_tokens": 34, "total_tokens": 34},
},
),
],
)
def test_the_sync_and_async_success_handlers_both_read_the_usage_from_one_relayed_body(path, request_data, body):
logging_obj = _relay_logging_obj(path, request_data)
relayed = _relayed_json(path, body)
sync_handler_view = logging_obj.normalize_logging_result(result=relayed)
async_handler_view = logging_obj.normalize_logging_result(result=relayed)
assert [view.usage.prompt_tokens for view in (sync_handler_view, async_handler_view)] == [34, 34]
@pytest.mark.asyncio
async def test_a_relayed_text_completion_stream_is_logged_with_the_usage_chunk_vllm_sent():
sse = _text_stream("cmpl-vllm-5", [(0, " Par"), (0, "is.")], USAGE)
logging_obj = await _log_stream("v1/completions", sse, {"model": MODEL, "prompt": PROMPT, "stream": True})
payload = logging_obj.model_call_details["standard_logging_object"]
assert payload["id"] == "cmpl-vllm-5"
assert payload["response"]["choices"][0]["text"] == " Paris."
assert (payload["prompt_tokens"], payload["completion_tokens"]) == (34, 4)
@pytest.mark.asyncio
async def test_a_relayed_text_completion_stream_without_a_usage_chunk_counts_every_prompt_and_choice():
prompts = ["Hi", "Yo"]
sse = _text_stream("cmpl-vllm-6", [(0, "Hello"), (1, "Hey"), (0, " there"), (1, " you")], None)
logging_obj = await _log_stream("completions", sse, {"model": MODEL, "prompt": prompts, "stream": True})
payload = logging_obj.model_call_details["standard_logging_object"]
assert payload["id"] == "cmpl-vllm-6"
assert [choice["text"] for choice in payload["response"]["choices"]] == ["Hello there", "Hey you"]
assert payload["prompt_tokens"] == sum(litellm.token_counter(model=MODEL, text=prompt) for prompt in prompts)
assert payload["completion_tokens"] == sum(
litellm.token_counter(model=MODEL, text=text, count_response_tokens=True) for text in ["Hello there", "Hey you"]
)

243
tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py Normal file
View file

@ -0,0 +1,243 @@
import json
from collections.abc import AsyncIterator, Iterable
import pytest
from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
stream_usage_relay,
without_injected_usage_frames,
)
MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
FINGERPRINT = "fp-vllm-test"
def _frame(payload: dict) -> bytes:
return b"data: " + json.dumps(payload).encode() + b"\n\n"
def _chunk(identity: str, **fields) -> dict:
return {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": MODEL, **fields}
def _content(index: int, **fields) -> bytes:
return _frame(
_chunk("chatcmpl-1", choices=[{"index": index, "delta": {"content": "Hi"}, "finish_reason": None}], **fields)
)
def _finish(index: int, **fields) -> bytes:
return _frame(_chunk("chatcmpl-1", choices=[{"index": index, "delta": {}, "finish_reason": "stop"}], **fields))
USAGE_COUNTS = {"prompt_tokens": 34, "completion_tokens": 4, "total_tokens": 38}
CONTENT = _content(0)
FINISH = _finish(0)
USAGE = _frame(_chunk("chatcmpl-1", choices=[], usage=USAGE_COUNTS, system_fingerprint=FINGERPRINT))
DONE = b"data: [DONE]\n\n"
async def _pieces(pieces: Iterable[bytes]) -> AsyncIterator[bytes]:
for piece in pieces:
yield piece
async def _relayed(pieces: Iterable[bytes]) -> bytes:
return b"".join([piece async for piece in without_injected_usage_frames(_pieces(pieces))])
def _payloads(relayed: bytes) -> list[dict]:
return [json.loads(line[len("data: ") :]) for line in relayed.decode().split("\n\n") if line.startswith("data: {")]
class TestStreamUsageRelay:
def test_a_chat_stream_nobody_asked_usage_for_gets_it_injected_and_stripped(self):
relay = stream_usage_relay(
endpoint="chat/completions",
request_body={"model": MODEL, "messages": [], "stream": True},
general_settings={},
supports_stream_options=lambda: True,
)
assert relay.body == {"model": MODEL, "messages": [], "stream": True, "stream_options": {"include_usage": True}}
assert relay.strips_usage is True
def test_a_chat_stream_the_caller_asked_usage_for_is_relayed_untouched(self):
body = {"model": MODEL, "messages": [], "stream": True, "stream_options": {"include_usage": True}}
relay = stream_usage_relay(
endpoint="v1/chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
)
assert relay.body == body
assert relay.strips_usage is False
def test_a_deployment_without_stream_options_support_is_left_alone(self):
body = {"model": MODEL, "messages": [], "stream": True}
relay = stream_usage_relay(
endpoint="chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: False
)
assert relay.body == body
assert relay.strips_usage is False
def test_a_text_completion_stream_gets_usage_injected_and_stripped_like_a_chat_stream(self):
body = {"model": MODEL, "prompt": "Hi", "stream": True}
by_default = stream_usage_relay(
endpoint="completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
)
always = stream_usage_relay(
endpoint="v1/completions",
request_body=body,
general_settings={"always_include_stream_usage": True},
supports_stream_options=lambda: True,
)
assert by_default.body == always.body == {**body, "stream_options": {"include_usage": True}}
assert (by_default.strips_usage, always.strips_usage) == (True, True)
def test_turning_stream_usage_off_relays_the_stream_untouched(self):
body = {"model": MODEL, "messages": [], "stream": True}
relay = stream_usage_relay(
endpoint="chat/completions",
request_body=body,
general_settings={"always_include_stream_usage": False},
supports_stream_options=lambda: True,
)
assert relay.body == body
assert relay.strips_usage is False
@pytest.mark.parametrize(
"stream_options",
[True, "include_usage", ["include_usage"], {"include_usage": "true"}, {"include_usage": 1}],
)
def test_stream_options_the_proxy_cannot_read_reach_the_server_as_the_caller_sent_them(self, stream_options):
body = {"model": MODEL, "messages": [], "stream": True, "stream_options": stream_options}
relay = stream_usage_relay(
endpoint="chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
)
assert relay.body == body
assert relay.strips_usage is False
def test_continuous_usage_stats_alone_is_not_turned_into_per_chunk_usage(self):
body = {"model": MODEL, "messages": [], "stream": True, "stream_options": {"continuous_usage_stats": True}}
relay = stream_usage_relay(
endpoint="chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
)
assert relay.body == {**body, "stream_options": {"include_usage": True}}
assert relay.strips_usage is True
def test_a_non_completion_endpoint_is_relayed_untouched(self):
body = {"model": MODEL, "input": "Hi", "stream": True}
relay = stream_usage_relay(
endpoint="embeddings", request_body=body, general_settings={}, supports_stream_options=lambda: True
)
assert relay.body == body
assert relay.strips_usage is False
class TestWithoutInjectedUsageFrames:
@pytest.mark.asyncio
async def test_the_usage_frame_is_dropped_and_its_fingerprint_moves_onto_the_finish_frame(self):
relayed = await _relayed([CONTENT, FINISH, USAGE, DONE])
payloads = _payloads(relayed)
assert [payload["choices"] for payload in payloads] == [
[{"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}],
[{"index": 0, "delta": {}, "finish_reason": "stop"}],
]
assert "usage" not in relayed.decode()
assert payloads[1]["system_fingerprint"] == FINGERPRINT
assert relayed.endswith(DONE)
@pytest.mark.asyncio
@pytest.mark.parametrize("frame_end", [b"\n\n", b"\r\n\r\n", b"\r\r"])
async def test_frames_split_across_upstream_reads_relay_the_same_bytes(self, frame_end):
whole = (CONTENT + FINISH + USAGE + DONE).replace(b"\n\n", frame_end)
assert await _relayed([whole[index : index + 1] for index in range(len(whole))]) == await _relayed([whole])
@pytest.mark.asyncio
async def test_a_stream_without_a_usage_frame_still_relays_its_finish_frame(self):
assert await _relayed([CONTENT, FINISH, DONE]) == CONTENT + FINISH + DONE
@pytest.mark.asyncio
async def test_a_finish_frame_already_fingerprinted_is_not_held(self):
stamped_finish = _frame(
_chunk(
"chatcmpl-1", choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], system_fingerprint="fp-0"
)
)
relayed = await _relayed([CONTENT, stamped_finish, USAGE, DONE])
assert relayed == CONTENT + stamped_finish + DONE
@pytest.mark.asyncio
async def test_every_finish_frame_gets_the_fingerprint_when_the_stream_has_several_choices(self):
relayed = await _relayed([_content(0), _content(1), _finish(0), _finish(1), USAGE, DONE])
payloads = _payloads(relayed)
finishes = {
payload["choices"][0]["index"]: payload.get("system_fingerprint")
for payload in payloads
if payload["choices"][0]["finish_reason"]
}
assert finishes == {0: FINGERPRINT, 1: FINGERPRINT}
assert "usage" not in relayed.decode()
@pytest.mark.asyncio
async def test_a_choice_that_finishes_while_another_still_streams_keeps_its_place(self):
last_token_finish = _frame(
_chunk("chatcmpl-1", choices=[{"index": 0, "delta": {"content": "!"}, "finish_reason": "stop"}])
)
upstream = [_content(0), _content(1), last_token_finish, _content(1), _content(1), _finish(1), USAGE, DONE]
relayed = await _relayed(upstream)
assert [payload["choices"] for payload in _payloads(relayed)] == [
payload["choices"] for payload in _payloads(b"".join(upstream[:-2]))
]
assert _payloads(relayed)[-1]["system_fingerprint"] == FINGERPRINT
@pytest.mark.asyncio
async def test_a_server_that_forces_usage_onto_every_chunk_is_relayed_untouched(self):
forced = [
_content(0, usage=USAGE_COUNTS),
_finish(0, usage=USAGE_COUNTS),
USAGE,
DONE,
]
assert await _relayed(forced) == b"".join(forced)
@pytest.mark.asyncio
async def test_closing_the_relay_early_closes_the_upstream(self):
closed = []
async def upstream() -> AsyncIterator[bytes]:
try:
yield CONTENT
yield FINISH
finally:
closed.append(True)
relay = without_injected_usage_frames(upstream())
assert await anext(relay) == CONTENT
await relay.aclose()
assert closed == [True]
@pytest.mark.asyncio
async def test_an_unterminated_tail_is_flushed_when_the_upstream_closes(self):
assert await _relayed([CONTENT, b'data: {"partial": true']) == CONTENT + b'data: {"partial": true'

171
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -2283,6 +2283,121 @@ class TestVLLMProxyRoute:
mock_is_router.assert_called_once()
mock_llm_router.allm_passthrough_route.assert_awaited_once()
@pytest.mark.asyncio
@patch( # test-quality-ok: patching litellm internal for unit test isolation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
return_value={"model": "router-model", "messages": [{"role": "user", "content": "Hi"}], "stream": True},
)
@patch( # test-quality-ok: patching litellm internal for unit test isolation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model",
return_value=True,
)
@patch(
"litellm.proxy.proxy_server.general_settings", {}
) # test-quality-ok: the proxy config surface for always_include_stream_usage
@patch(
"litellm.proxy.proxy_server.llm_router"
) # test-quality-ok: patching litellm internal for unit test isolation
async def test_vllm_proxy_route_asks_the_router_model_for_stream_usage_and_hides_it_from_the_client(
self, mock_llm_router, mock_is_router, mock_get_body
):
chunk = {
"id": "chatcmpl-1",
"object": "chat.completion.chunk",
"created": 1,
"model": "Qwen/Qwen2.5-0.5B-Instruct",
}
frames = [
{**chunk, "choices": [{"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}]},
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
{
**chunk,
"choices": [],
"usage": {"prompt_tokens": 34, "completion_tokens": 4},
"system_fingerprint": "fp-1",
},
]
upstream_body = (
b"".join(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames) + b"data: [DONE]\n\n"
)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.headers = {"content-type": "application/json"}
mock_request.query_params = {}
mock_llm_router.resolved_litellm_models = MagicMock(return_value=("hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct",))
mock_llm_router.allm_passthrough_route = AsyncMock(
return_value=httpx.Response(200, content=upstream_body, headers={"content-type": "text/event-stream"})
)
result = await vllm_proxy_route(
endpoint="/chat/completions",
request=mock_request,
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"),
)
relayed_body = mock_llm_router.allm_passthrough_route.await_args.kwargs["json"]
assert relayed_body["stream_options"] == {"include_usage": True}
assert isinstance(result, StreamingResponse)
body = b"".join([piece async for piece in result.body_iterator])
payloads = [
json.loads(line[len("data: ") :]) for line in body.decode().split("\n\n") if line.startswith("data: {")
]
assert [payload["choices"] for payload in payloads] == [frames[0]["choices"], frames[1]["choices"]]
assert "usage" not in body.decode()
assert payloads[1]["system_fingerprint"] == "fp-1"
assert body.endswith(b"data: [DONE]\n\n")
@pytest.mark.asyncio
@pytest.mark.parametrize("stream_options", [None, {"include_usage": True}])
@patch( # test-quality-ok: patching litellm internal for unit test isolation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model",
return_value=True,
)
@patch(
"litellm.proxy.proxy_server.general_settings", {}
) # test-quality-ok: the proxy config surface for always_include_stream_usage
@patch(
"litellm.proxy.proxy_server.llm_router"
) # test-quality-ok: patching litellm internal for unit test isolation
async def test_vllm_proxy_route_closes_the_upstream_stream_when_the_client_disconnects(
self, mock_llm_router, mock_is_router, stream_options
):
closed = []
class UpstreamStream(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'data: {"id":"chatcmpl-1","choices":[]}\n\n'
yield b"data: [DONE]\n\n"
async def aclose(self) -> None:
closed.append(True)
body = {"model": "router-model", "messages": [{"role": "user", "content": "Hi"}], "stream": True}
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.headers = {"content-type": "application/json"}
mock_request.query_params = {}
mock_llm_router.resolved_litellm_models = MagicMock(return_value=("hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct",))
mock_llm_router.allm_passthrough_route = AsyncMock(
return_value=httpx.Response(200, stream=UpstreamStream(), headers={"content-type": "text/event-stream"})
)
with patch( # test-quality-ok: patching litellm internal for unit test isolation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
return_value={**body, **({"stream_options": stream_options} if stream_options else {})},
):
result = await vllm_proxy_route(
endpoint="/chat/completions",
request=mock_request,
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"),
)
await anext(result.body_iterator)
await result.body_iterator.aclose()
assert closed == [True]
@pytest.mark.asyncio
@patch( # test-quality-ok: patching litellm internal for unit test isolation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
@ -4472,7 +4587,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
(b"authorization", f"Bearer {self.MASTER_KEY}".encode()),
(b"content-type", b"application/json"),
],
authenticated=UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN),
authenticated=UserAPIKeyAuth(
api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN
),
)
assert forwarded is None, "the master key must never reach the upstream forwarder"
assert raised is not None and raised.status_code == 401
@ -7752,13 +7869,17 @@ class TestTypeSafePassthroughRoute:
monkeypatch.setenv("OPENROUTER_API_BASE", "https://typesafe.example/base")
monkeypatch.setenv("LAYA_API_BASE", "https://typesafe.example/base")
monkeypatch.setenv("BESPOKE_API_BASE", "https://typesafe.example/base")
model: Final = {"typesafe": "jev-latest", "laya": "english", "bespoke": "nimble-latest"}.get(provider, "test-generative-model")
model: Final = {"typesafe": "jev-latest", "laya": "english", "bespoke": "nimble-latest"}.get(
provider, "test-generative-model"
)
permission_model: Final = f"{provider}/{model}" if provider in ("laya", "bespoke") else model
auth: Final = UserAPIKeyAuth(
api_key="sk-limited",
tpm_limit=token_limit if quota_scope == "key" else None,
project_id="test-project" if quota_scope == "project_output" else None,
project_metadata={"model_otpm_limit": {permission_model: token_limit}} if quota_scope == "project_output" else {},
project_metadata={"model_otpm_limit": {permission_model: token_limit}}
if quota_scope == "project_output"
else {},
)
monkeypatch.setitem(proxy_server.app.dependency_overrides, user_api_key_auth, lambda: auth)
body: Final = (
@ -7780,7 +7901,9 @@ class TestTypeSafePassthroughRoute:
return httpx.Response(200, json={"model": model})
with respx.mock(assert_all_called=False) as upstream:
route: Final = upstream.post(f"https://typesafe.example/base/v1/{endpoint}").mock(side_effect=upstream_response)
route: Final = upstream.post(f"https://typesafe.example/base/v1/{endpoint}").mock(
side_effect=upstream_response
)
response: Final = client.post(f"/{provider}/v1/{endpoint}", json=body)
assert response.status_code == (429 if token_limit == 0 else 200), response.text
@ -7856,11 +7979,14 @@ class TestOssDecisionPassthroughRoute:
"questions": {"department": {"type": "choice", "criteria": {"billing": "refunds"}}},
}
answer: Final = {
"model": "laya-rl-agent" if provider == "laya" else checkpoint, "answers": {},
"model": "laya-rl-agent" if provider == "laya" else checkpoint,
"answers": {},
**({"routing": {"model": checkpoint}} if provider == "laya" else {}),
}
with respx.mock(assert_all_called=True) as upstream:
route: Final = upstream.post(f"http://{provider}.test/base/v1/systemone?trace=yes").respond(200, json=answer)
route: Final = upstream.post(f"http://{provider}.test/base/v1/systemone?trace=yes").respond(
200, json=answer
)
response: Final = client.post(
f"/{provider}/v1/systemone?trace=yes",
json=body,
@ -7882,14 +8008,18 @@ class TestOssDecisionPassthroughRoute:
assert f"{provider.upper()}_API_BASE" in response.text
assert len(upstream.calls) == 0
def test_oss_does_not_forward_unsupported_endpoints(self, client: TestClient, provider: str, checkpoint: str) -> None:
def test_oss_does_not_forward_unsupported_endpoints(
self, client: TestClient, provider: str, checkpoint: str
) -> None:
with respx.mock(assert_all_called=False) as upstream:
response: Final = client.post(f"/{provider}/v1/evaluate", json={"model": checkpoint})
assert response.status_code == 404
assert len(upstream.calls) == 0
@pytest.mark.parametrize("model", [None, "auto", "jev-latest"])
def test_oss_rejects_implicit_checkpoint_selection(self, client: TestClient, model: str | None, provider: str) -> None:
def test_oss_rejects_implicit_checkpoint_selection(
self, client: TestClient, model: str | None, provider: str
) -> None:
with respx.mock(assert_all_called=False) as upstream:
response: Final = client.post(f"/{provider}/v1/systemone", json={"model": model})
assert response.status_code == 400
@ -7908,7 +8038,6 @@ class TestOssDecisionPassthroughRoute:
assert response.status_code == 400
assert not route.called
@pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"])
def test_oss_hooks_enforce_canonical_model_limits_and_keep_native_wire_body(
self, client: TestClient, monkeypatch: pytest.MonkeyPatch, metadata_slot: str, provider: str, checkpoint: str
@ -7921,8 +8050,10 @@ class TestOssDecisionPassthroughRoute:
cache: Final = DualCache()
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(cache))
auth: Final = UserAPIKeyAuth(
api_key="oss-native-rpm", metadata={"model_rpm_limit": {f"{provider}/{checkpoint}": 1}},
api_key="oss-native-rpm",
metadata={"model_rpm_limit": {f"{provider}/{checkpoint}": 1}},
)
def authenticated_key() -> UserAPIKeyAuth:
return auth
@ -7930,8 +8061,11 @@ class TestOssDecisionPassthroughRoute:
class LimitHook(CustomLogger):
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache,
data: dict[str, object], call_type: CallTypesLiteral,
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
assert data["model"] == f"{provider}/{checkpoint}"
metadata: Final = data.get(metadata_slot)
@ -7943,9 +8077,11 @@ class TestOssDecisionPassthroughRoute:
monkeypatch.setattr(litellm, "callbacks", [LimitHook()])
body: Final = {
"model": checkpoint, "state": "refund",
"model": checkpoint,
"state": "refund",
metadata_slot: {
"customer_label": "retained", "model_group": "unbounded-client-choice",
"customer_label": "retained",
"model_group": "unbounded-client-choice",
"standard_logging_guardrail_information": [{"guardrail_cost": 25.0}],
},
}
@ -7967,8 +8103,11 @@ class TestOssDecisionPassthroughRoute:
class CheckpointHook(CustomLogger):
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache,
data: dict[str, object], call_type: CallTypesLiteral,
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
assert data["model"] == f"{provider}/{checkpoint}"
return {**data, "model": f"{provider}/{changed_checkpoint}"}

8
tests/unit/proxy/test_common_request_processing.py
View file

@ -8282,10 +8282,10 @@ def _apply_stream_usage_tracking(
route_type: str,
supports_stream_options: Callable[[], bool] = lambda: True,
) -> None:
from litellm.proxy.common_request_processing import _stream_usage_tracking_updates
from litellm.proxy.common_request_processing import stream_usage_tracking_updates
data.update(
_stream_usage_tracking_updates(
stream_usage_tracking_updates(
data=data,
general_settings=general_settings,
route_type=route_type,
@ -8430,10 +8430,10 @@ class TestApplyStreamUsageTracking:
class TestModelDeploymentsSupportStreamOptions:
def _support(self, model, llm_router=None, team_id=None) -> bool:
from litellm.proxy.common_request_processing import (
_model_deployments_support_stream_options,
model_deployments_support_stream_options,
)
return _model_deployments_support_stream_options(model=model, llm_router=llm_router, team_id=team_id)
return model_deployments_support_stream_options(model=model, llm_router=llm_router, team_id=team_id)
def test_openai_compatible_deployment_supports_stream_options(self):
router = litellm.Router(