diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index b020122ee3b..f521b2c8edd 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 944d3c57f7d..9b10b95956f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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: diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 501122783eb..c3d5540aaea 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -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"(? 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) diff --git a/litellm/llms/openai_like/passthrough/transformation.py b/litellm/llms/openai_like/passthrough/transformation.py new file mode 100644 index 00000000000..df1cebc8519 --- /dev/null +++ b/litellm/llms/openai_like/passthrough/transformation.py @@ -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 diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py index 71e5923323e..e7da6e54410 100644 --- a/litellm/llms/vllm/passthrough/transformation.py +++ b/litellm/llms/vllm/passthrough/transformation.py @@ -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) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7ed486922ca..58e8866ce99 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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, diff --git a/litellm/proxy/pass_through_endpoints/injected_stream_usage.py b/litellm/proxy/pass_through_endpoints/injected_stream_usage.py new file mode 100644 index 00000000000..46a2fc03f73 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/injected_stream_usage.py @@ -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() diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1eefb58d7fa..70cee344c91 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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(), diff --git a/tests/integration/routing/test_azure_passthrough_spend_row.py b/tests/integration/routing/test_azure_passthrough_spend_row.py new file mode 100644 index 00000000000..c2a33aadadb --- /dev/null +++ b/tests/integration/routing/test_azure_passthrough_spend_row.py @@ -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 diff --git a/tests/integration/routing/test_vllm_passthrough_spend_row.py b/tests/integration/routing/test_vllm_passthrough_spend_row.py new file mode 100644 index 00000000000..ac42b84102b --- /dev/null +++ b/tests/integration/routing/test_vllm_passthrough_spend_row.py @@ -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 diff --git a/tests/integration/routing/test_vllm_passthrough_stream_relay.py b/tests/integration/routing/test_vllm_passthrough_stream_relay.py new file mode 100644 index 00000000000..6ab7469aa63 --- /dev/null +++ b/tests/integration/routing/test_vllm_passthrough_stream_relay.py @@ -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 diff --git a/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py b/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py new file mode 100644 index 00000000000..02e70a82c1c --- /dev/null +++ b/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py @@ -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 diff --git a/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py index bc8c1000fa4..f608e91702c 100644 --- a/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -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], diff --git a/tests/unit/llms/vllm/__init__.py b/tests/unit/llms/vllm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vllm/passthrough/__init__.py b/tests/unit/llms/vllm/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py b/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py new file mode 100644 index 00000000000..bc8b7111fbb --- /dev/null +++ b/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py @@ -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"] + ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py b/tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py new file mode 100644 index 00000000000..82c7ce8b864 --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py @@ -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' diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index d25b4470668..e2ecfbb8ba9 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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}"} diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 85a62614448..09ed6399d68 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -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(