mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
fix(hosted_vllm): keep vLLM choice fields for null, cache hits, n>1, and streamed completions
Response conversion keeps stop_reason and token_ids on the choice for any OpenAI compatible response, cache replays carry them, finish fields are recorded per choice index, text completion streams keep token_ids, and the proxy passes only these two keys through, nulls included.
This commit is contained in:
mateo-berri
parent
a7ac7ad54e
commit
efefc437ce
13 changed files with 272 additions and 94 deletions
|
|
@ -688,6 +688,7 @@ STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping
|
|||
STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8")
|
||||
STREAM_FINISH_CHOICE_FIELDS: Final = ("stop_reason",)
|
||||
STREAM_CHUNK_CHOICE_FIELDS: Final = ("token_ids",)
|
||||
VLLM_CHOICE_FIELDS: Final = STREAM_FINISH_CHOICE_FIELDS + STREAM_CHUNK_CHOICE_FIELDS
|
||||
### SPEND TRACKING ###
|
||||
DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float(
|
||||
"DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ import json
|
|||
import re
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
|
|
@ -11,7 +12,12 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.constants import (
|
||||
RESPONSE_FORMAT_TOOL_NAME,
|
||||
STREAM_CHUNK_CHOICE_FIELDS,
|
||||
STREAM_FINISH_CHOICE_FIELDS,
|
||||
VLLM_CHOICE_FIELDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_reasoning_content,
|
||||
|
|
@ -154,8 +160,17 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
|
|||
# later can't silently re-introduce the duplication.
|
||||
choice.delta = Delta(content=choice.delta.content)
|
||||
choice.logprobs = None
|
||||
if hasattr(choice, "enhancements"):
|
||||
del choice.enhancements
|
||||
_drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS))
|
||||
|
||||
|
||||
def _drop_choice_fields(choice: StreamingChoices, fields: Iterable[str]) -> None:
|
||||
for field in fields:
|
||||
if hasattr(choice, field):
|
||||
delattr(choice, field)
|
||||
|
||||
|
||||
def _vllm_choice_fields(choice: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return MappingProxyType({field: choice[field] for field in VLLM_CHOICE_FIELDS if field in choice})
|
||||
|
||||
|
||||
def _first_replay_delta(
|
||||
|
|
@ -242,7 +257,9 @@ async def convert_to_streaming_response_async(
|
|||
|
||||
logprobs = choice.get("logprobs", None)
|
||||
|
||||
choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs)
|
||||
choice = StreamingChoices(
|
||||
finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs, **_vllm_choice_fields(choice)
|
||||
)
|
||||
choice_list.append(choice)
|
||||
|
||||
model_response_object.choices = choice_list
|
||||
|
|
@ -292,6 +309,8 @@ async def convert_to_streaming_response_async(
|
|||
if i > 0:
|
||||
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
|
||||
slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
|
||||
if i < last_idx:
|
||||
_drop_choice_fields(slice_chunk.choices[0], STREAM_FINISH_CHOICE_FIELDS)
|
||||
if i == last_idx and original_usage is not None:
|
||||
setattr(slice_chunk, "usage", original_usage)
|
||||
yield slice_chunk
|
||||
|
|
@ -332,6 +351,7 @@ def convert_to_streaming_response(
|
|||
delta=delta,
|
||||
logprobs=logprobs,
|
||||
enhancements=enhancements,
|
||||
**_vllm_choice_fields(choice),
|
||||
)
|
||||
|
||||
choice_list.append(choice)
|
||||
|
|
@ -377,6 +397,8 @@ def convert_to_streaming_response(
|
|||
if i > 0:
|
||||
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
|
||||
slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
|
||||
if i < last_idx:
|
||||
_drop_choice_fields(slice_chunk.choices[0], STREAM_FINISH_CHOICE_FIELDS)
|
||||
if i == last_idx and original_usage is not None:
|
||||
setattr(slice_chunk, "usage", original_usage)
|
||||
yield slice_chunk
|
||||
|
|
@ -834,6 +856,7 @@ def convert_to_model_response_object(
|
|||
logprobs=logprobs,
|
||||
enhancements=enhancements,
|
||||
provider_specific_fields=provider_specific_fields,
|
||||
**_vllm_choice_fields(choice),
|
||||
)
|
||||
if provider_sent_logprobs and logprobs is None:
|
||||
choice.logprobs = None
|
||||
|
|
|
|||
|
|
@ -115,6 +115,7 @@ def _prepend_upstream_choice_extras(
|
|||
|
||||
|
||||
_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_NO_FINISH_CHOICE_FIELDS: Final[Mapping[int, Mapping[str, object]]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None:
|
||||
|
|
@ -139,8 +140,12 @@ def _upstream_choice_index(upstream_chunk: object) -> int:
|
|||
|
||||
|
||||
def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
|
||||
first_choice: Final = _upstream_first_choice(upstream_chunk)
|
||||
upstream_extras: Final = first_choice.model_extra if first_choice is not None else None
|
||||
if not upstream_extras:
|
||||
return _NO_UPSTREAM_CHOICE_FIELDS
|
||||
return MappingProxyType(
|
||||
{field: value for field, value in _upstream_choice_fields(upstream_chunk, fields).items() if value is not None}
|
||||
{field: upstream_extras[field] for field in fields if upstream_extras.get(field) is not None}
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -373,7 +378,7 @@ class CustomStreamWrapper:
|
|||
self.system_fingerprint: str | None = None
|
||||
self._provider_response_model: str | None = None
|
||||
self.received_finish_reason: str | None = None
|
||||
self.received_finish_choice_fields: Mapping[str, object] = _NO_UPSTREAM_CHOICE_FIELDS
|
||||
self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_FINISH_CHOICE_FIELDS
|
||||
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
|
||||
self.special_tokens = [
|
||||
"<|assistant|>",
|
||||
|
|
@ -818,6 +823,7 @@ class CustomStreamWrapper:
|
|||
"logprobs": logprobs,
|
||||
"original_chunk": str_line,
|
||||
"usage": usage,
|
||||
"chunk_choice_fields": _non_null_upstream_choice_fields(str_line, STREAM_CHUNK_CHOICE_FIELDS),
|
||||
}
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
|
@ -842,7 +848,7 @@ class CustomStreamWrapper:
|
|||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def handle_openai_text_completion_chunk(self, chunk):
|
||||
def handle_openai_text_completion_chunk(self, chunk: object):
|
||||
try:
|
||||
text = ""
|
||||
is_finished = False
|
||||
|
|
@ -860,6 +866,7 @@ class CustomStreamWrapper:
|
|||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
"usage": usage,
|
||||
"chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -1170,6 +1177,9 @@ class CustomStreamWrapper:
|
|||
)
|
||||
|
||||
is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj)
|
||||
chunk_choice_fields: Final[Mapping[str, object]] = response_obj.get(
|
||||
"chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS
|
||||
)
|
||||
|
||||
if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string
|
||||
self.raise_on_model_repetition()
|
||||
|
|
@ -1262,6 +1272,7 @@ class CustomStreamWrapper:
|
|||
if _index is not None:
|
||||
model_response.choices[0].index = _index
|
||||
|
||||
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
|
||||
self._optional_combine_thinking_block_in_choices(model_response=model_response)
|
||||
|
||||
return model_response
|
||||
|
|
@ -1305,25 +1316,20 @@ class CustomStreamWrapper:
|
|||
model_response.choices[0].finish_reason = map_finish_reason(
|
||||
finish_reason=self.received_finish_reason
|
||||
) # ensure consistent output to openai
|
||||
_apply_choice_fields(model_response.choices[0], self.received_finish_choice_fields)
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
if response_obj.get("is_finished"):
|
||||
_apply_choice_fields(
|
||||
model_response.choices[0],
|
||||
_non_null_upstream_choice_fields(_original_chunk, STREAM_CHUNK_CHOICE_FIELDS),
|
||||
)
|
||||
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
|
||||
|
||||
self.sent_last_chunk = True
|
||||
|
||||
return model_response
|
||||
elif self._has_special_delta_content(model_response):
|
||||
return self._handle_special_delta_content(model_response)
|
||||
elif upstream_chunk_fields := _non_null_upstream_choice_fields(
|
||||
response_obj.get("original_chunk"), STREAM_CHUNK_CHOICE_FIELDS
|
||||
):
|
||||
elif chunk_choice_fields:
|
||||
upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
|
||||
model_response.choices[0].delta = Delta(content=None) # rebind-ok: chunks are built in place here
|
||||
model_response.choices[0].index = upstream_index # rebind-ok: chunks are built in place here
|
||||
_apply_choice_fields(model_response.choices[0], upstream_chunk_fields)
|
||||
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
|
||||
return model_response
|
||||
else:
|
||||
if hasattr(model_response, "usage"):
|
||||
|
|
@ -1543,7 +1549,7 @@ class CustomStreamWrapper:
|
|||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
self.received_finish_choice_fields = _upstream_choice_fields(chunk, STREAM_FINISH_CHOICE_FIELDS)
|
||||
self._record_finish_choice_fields(chunk)
|
||||
if response_obj["usage"] is not None:
|
||||
_text_completion_usage: Final[Usage] = response_obj["usage"]
|
||||
setattr(
|
||||
|
|
@ -1606,6 +1612,7 @@ class CustomStreamWrapper:
|
|||
self.system_fingerprint = cached_chunk.system_fingerprint
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
self._record_finish_choice_fields(cached_chunk)
|
||||
else: # openai / azure chat model
|
||||
if isinstance(chunk, ChatCompletionChunk):
|
||||
self._openai_handler_receives_chat_chunks = True
|
||||
|
|
@ -1633,7 +1640,7 @@ class CustomStreamWrapper:
|
|||
f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}"
|
||||
)
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
self.received_finish_choice_fields = _upstream_choice_fields(chunk, STREAM_FINISH_CHOICE_FIELDS)
|
||||
self._record_finish_choice_fields(chunk)
|
||||
if response_obj.get("original_chunk", None) is not None:
|
||||
if hasattr(response_obj["original_chunk"], "id"):
|
||||
model_response = self.set_model_id(response_obj["original_chunk"].id, model_response)
|
||||
|
|
@ -1970,6 +1977,19 @@ class CustomStreamWrapper:
|
|||
if self.logging_obj.is_sync_litellm_request(litellm_params):
|
||||
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
|
||||
|
||||
def _record_finish_choice_fields(self, upstream_chunk: object) -> None:
|
||||
self.received_finish_choice_fields = MappingProxyType(
|
||||
{
|
||||
**self.received_finish_choice_fields,
|
||||
_upstream_choice_index(upstream_chunk): _upstream_choice_fields(
|
||||
upstream_chunk, STREAM_FINISH_CHOICE_FIELDS
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None:
|
||||
_apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS))
|
||||
|
||||
def finish_reason_handler(self):
|
||||
model_response: Final = self.model_response_creator(
|
||||
hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id}
|
||||
|
|
@ -1979,7 +1999,7 @@ class CustomStreamWrapper:
|
|||
model_response.choices[0].finish_reason = _finish_reason
|
||||
else:
|
||||
model_response.choices[0].finish_reason = "stop"
|
||||
_apply_choice_fields(model_response.choices[0], self.received_finish_choice_fields)
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
|
||||
## if tool use
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -3,11 +3,8 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions`
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
from collections.abc import Coroutine
|
||||
from typing import Final, Literal, cast, overload
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_image_mime_type_from_url,
|
||||
|
|
@ -25,24 +22,10 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionVideoObject,
|
||||
ChatCompletionVideoUrlObject,
|
||||
)
|
||||
from litellm.types.utils import Choices, ModelResponse
|
||||
|
||||
from ....utils import remove_additional_properties, remove_strict_from_schema
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
_VLLM_CHOICE_FIELDS: Final = ("stop_reason", "token_ids")
|
||||
|
||||
|
||||
def _vllm_choice_fields(choice: Choices) -> Mapping[str, object]:
|
||||
provider_fields: Final[Mapping[str, object]] = (
|
||||
choice.provider_specific_fields or {} if "provider_specific_fields" in choice else {}
|
||||
)
|
||||
return {field: provider_fields[field] for field in _VLLM_CHOICE_FIELDS if field in provider_fields}
|
||||
|
||||
|
||||
class HostedVLLMChatConfig(OpenAIGPTConfig):
|
||||
def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]:
|
||||
|
|
@ -123,38 +106,6 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
|
|||
|
||||
return super().map_openai_params(non_default_params, optional_params, model, drop_params)
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ModelResponse,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature
|
||||
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
|
||||
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
|
||||
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
|
||||
encoding: "Tokenizer | None",
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ModelResponse:
|
||||
response: Final = super().transform_response( # pyright: ignore[reportUnknownMemberType] # OpenAIGPTConfig.transform_response takes bare dicts
|
||||
model=model,
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data=request_data,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=encoding,
|
||||
api_key=api_key,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
for choice in response.choices:
|
||||
for field, value in _vllm_choice_fields(choice).items():
|
||||
choice[field] = value
|
||||
return response
|
||||
|
||||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: str | None, api_key: str | None
|
||||
) -> tuple[str | None, str | None]:
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ from litellm.constants import (
|
|||
LITELLM_UI_SESSION_DURATION,
|
||||
OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS,
|
||||
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
|
||||
VLLM_CHOICE_FIELDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
|
|
@ -10065,7 +10066,9 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -
|
|||
return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending
|
||||
|
||||
|
||||
_STANDARD_STREAM_CHOICE_FIELDS: Final = frozenset({"finish_reason", "index", "delta", "logprobs", "enhancements"})
|
||||
def _vllm_choice_fields(choice: BaseModel) -> Mapping[str, object]:
|
||||
choice_extras: Final = choice.model_extra or {}
|
||||
return MappingProxyType({field: choice_extras[field] for field in VLLM_CHOICE_FIELDS if field in choice_extras})
|
||||
|
||||
|
||||
def _fast_serialize_simple_model_response_stream(
|
||||
|
|
@ -10088,9 +10091,11 @@ def _fast_serialize_simple_model_response_stream(
|
|||
return None
|
||||
|
||||
choice: Final = choices[0]
|
||||
if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None:
|
||||
return None
|
||||
if not isinstance(choice, BaseModel) or (choice.model_extra or {}).keys() - _STANDARD_STREAM_CHOICE_FIELDS:
|
||||
if (
|
||||
not isinstance(choice, BaseModel)
|
||||
or getattr(choice, "logprobs", None) is not None
|
||||
or getattr(choice, "enhancements", None) is not None
|
||||
):
|
||||
return None
|
||||
|
||||
delta: Final[object] = getattr(choice, "delta", None)
|
||||
|
|
@ -10119,7 +10124,7 @@ def _fast_serialize_simple_model_response_stream(
|
|||
if content is not None:
|
||||
delta_dict["content"] = content
|
||||
|
||||
choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict}
|
||||
choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict, **_vllm_choice_fields(choice)}
|
||||
finish_reason: Final = getattr(choice, "finish_reason", None)
|
||||
if finish_reason is not None:
|
||||
choice_dict["finish_reason"] = finish_reason
|
||||
|
|
@ -10170,10 +10175,7 @@ def _is_empty_streaming_choice(choice: StreamingChoices) -> bool:
|
|||
return False
|
||||
if getattr(choice, "logprobs", None) is not None:
|
||||
return False
|
||||
provider_choice_values: Final = (
|
||||
value for field, value in (choice.model_extra or {}).items() if field not in _STANDARD_STREAM_CHOICE_FIELDS
|
||||
)
|
||||
if not _all_values_none(provider_choice_values):
|
||||
if not _all_values_none(_vllm_choice_fields(choice).values()):
|
||||
return False
|
||||
delta: Final = choice.delta if hasattr(choice, "delta") else None
|
||||
if delta is None:
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ from litellm.constants import (
|
|||
SPEND_LOG_QUEUE_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
|
||||
VLLM_CHOICE_FIELDS,
|
||||
)
|
||||
from litellm.cost_calculator import BaseTokenUsageProcessor
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
|
|
@ -9917,7 +9918,7 @@ _PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [
|
|||
("message", "refusal"),
|
||||
("delta", "content"), # null in streaming chunks
|
||||
]
|
||||
_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs",)
|
||||
_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs", *VLLM_CHOICE_FIELDS)
|
||||
_PRESERVED_NONE_RESPONSE_FIELDS: Final = ("system_fingerprint",)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.constants import STREAM_FINISH_CHOICE_FIELDS
|
||||
from litellm.constants import VLLM_CHOICE_FIELDS
|
||||
from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse
|
||||
|
||||
|
||||
|
|
@ -41,7 +41,7 @@ class TextCompletionStreamWrapper:
|
|||
text_choices["reasoning_content"] = delta.get("reasoning_content")
|
||||
text_choices["index"] = chat_choice["index"]
|
||||
text_choices["finish_reason"] = chat_choice["finish_reason"]
|
||||
for field in STREAM_FINISH_CHOICE_FIELDS:
|
||||
for field in VLLM_CHOICE_FIELDS:
|
||||
if field in chat_choice:
|
||||
text_choices[field] = chat_choice[field]
|
||||
response["choices"] = [text_choices]
|
||||
|
|
|
|||
|
|
@ -2968,3 +2968,39 @@ def test_convert_to_model_response_object_function_output():
|
|||
assert result.usage.completion_tokens_details == CompletionTokensDetailsWrapper(reasoning_tokens=0)
|
||||
|
||||
assert result._hidden_params is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vllm_choice_fields",
|
||||
[{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": 151645}, {"stop_reason": None, "token_ids": None}],
|
||||
ids=["stop-string-with-token-ids", "stop-token", "natural-stop"],
|
||||
)
|
||||
def test_vllm_choice_fields_stay_on_the_choice_through_a_cache_round_trip(
|
||||
vllm_choice_fields: dict[str, object],
|
||||
) -> None:
|
||||
upstream: Final = {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "1, 2, 3, 4"},
|
||||
"finish_reason": "stop",
|
||||
**vllm_choice_fields,
|
||||
}
|
||||
],
|
||||
}
|
||||
response: Final = convert_to_model_response_object(response_object=upstream, model_response_object=ModelResponse())
|
||||
assert isinstance(response, ModelResponse)
|
||||
cache_hit: Final = convert_to_model_response_object(
|
||||
response_object=json.loads(response.model_dump_json()), model_response_object=ModelResponse()
|
||||
)
|
||||
assert isinstance(cache_hit, ModelResponse)
|
||||
|
||||
dumped_choices: Final = [json.loads(rebuilt.model_dump_json())["choices"][0] for rebuilt in (response, cache_hit)]
|
||||
assert [{field: choice.get(field, "missing") for field in vllm_choice_fields} for choice in dumped_choices] == [
|
||||
vllm_choice_fields,
|
||||
vllm_choice_fields,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -25,6 +25,10 @@ from litellm.exceptions import MidStreamFallbackError
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
convert_to_streaming_response,
|
||||
convert_to_streaming_response_async,
|
||||
)
|
||||
from litellm.litellm_core_utils.streaming_handler import (
|
||||
CustomStreamWrapper,
|
||||
_ProviderChunkEarlyReturn,
|
||||
|
|
@ -6048,6 +6052,13 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj:
|
|||
assert result.choices[0].delta.content == "hello"
|
||||
|
||||
|
||||
_TOKEN_IDS: Final = TypeAdapter(list[int] | None)
|
||||
|
||||
|
||||
def _joined_token_ids(choices: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
|
||||
return tuple(chain.from_iterable(_TOKEN_IDS.validate_python(choice.get("token_ids")) or () for choice in choices))
|
||||
|
||||
|
||||
def _hosted_vllm_stream_choices(
|
||||
logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]]
|
||||
) -> tuple[dict[str, object], ...]:
|
||||
|
|
@ -6091,7 +6102,7 @@ def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id(
|
|||
)
|
||||
|
||||
assert (emitted[-1]["finish_reason"], emitted[-1].get("stop_reason")) == ("stop", stop_reason)
|
||||
assert list(chain.from_iterable(choice.get("token_ids") or () for choice in emitted)) == [16, 11, 17, 20]
|
||||
assert _joined_token_ids(emitted) == (16, 11, 17, 20)
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None:
|
||||
|
|
@ -6131,6 +6142,71 @@ def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without
|
|||
]
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_finish_chunk_keeps_the_first_choice_stop_reason(logging_obj: Logging) -> None:
|
||||
emitted: Final = _hosted_vllm_stream_choices(
|
||||
logging_obj,
|
||||
(
|
||||
{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},
|
||||
{"index": 1, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},
|
||||
{"index": 0, "delta": {"content": "1, 2"}, "finish_reason": "stop", "stop_reason": ", 3"},
|
||||
{"index": 1, "delta": {"content": "1, 2, 3"}, "finish_reason": "stop", "stop_reason": 151645},
|
||||
),
|
||||
)
|
||||
|
||||
assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [", 3"]
|
||||
|
||||
|
||||
_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19)
|
||||
|
||||
|
||||
def _cached_vllm_completion() -> dict[str, object]:
|
||||
return {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "1, 2, 3, 4"},
|
||||
"finish_reason": "stop",
|
||||
"stop_reason": ", 5",
|
||||
"token_ids": list(_CACHED_VLLM_TOKEN_IDS),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
async def _replayed_cache_hit_choices(logging_obj: Logging, replay: str) -> tuple[dict[str, object], ...]:
|
||||
if replay == "async":
|
||||
async_wrapper: Final = CustomStreamWrapper(
|
||||
completion_stream=convert_to_streaming_response_async(_cached_vllm_completion()),
|
||||
model="m",
|
||||
custom_llm_provider="cached_response",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return tuple([chunk.choices[0].model_dump() async for chunk in async_wrapper if chunk.choices])
|
||||
sync_wrapper: Final = CustomStreamWrapper(
|
||||
completion_stream=convert_to_streaming_response(_cached_vllm_completion()),
|
||||
model="m",
|
||||
custom_llm_provider="cached_response",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return tuple(chunk.choices[0].model_dump() for chunk in sync_wrapper if chunk.choices)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("replay", ["sync", "async"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_id(
|
||||
logging_obj: Logging, replay: str
|
||||
) -> None:
|
||||
emitted: Final = await _replayed_cache_hit_choices(logging_obj, replay)
|
||||
|
||||
assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [", 5"]
|
||||
assert {choice.get("stop_reason") for choice in emitted[:-2]} == {None}
|
||||
assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS
|
||||
|
||||
|
||||
def _later_chunk_outputs(
|
||||
custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging
|
||||
) -> tuple[ModelResponseStream, ...]:
|
||||
|
|
|
|||
|
|
@ -576,13 +576,27 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes():
|
|||
}
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_keeps_provider_choice_fields():
|
||||
@pytest.mark.parametrize(
|
||||
("choice_extras", "serialized_extras"),
|
||||
[
|
||||
({"stop_reason": ",5", "token_ids": [20]}, {"stop_reason": ",5", "token_ids": [20]}),
|
||||
({"text": "1,2,3,4"}, {}),
|
||||
(
|
||||
{"finish_reason": "stop", "stop_reason": None, "text": "1,2,3,4"},
|
||||
{"finish_reason": "stop", "stop_reason": None},
|
||||
),
|
||||
],
|
||||
ids=["vllm-fields", "unrelated-extra", "natural-stop"],
|
||||
)
|
||||
def test_serialize_streaming_chunk_keeps_only_the_vllm_choice_fields(
|
||||
choice_extras: dict[str, object], serialized_extras: dict[str, object]
|
||||
):
|
||||
chunk: Final = _simple_chunk(content="1,2,3,4")
|
||||
chunk.choices[0]["stop_reason"] = ",5"
|
||||
chunk.choices[0]["token_ids"] = [20]
|
||||
for field, value in choice_extras.items():
|
||||
chunk.choices[0][field] = value
|
||||
|
||||
assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [
|
||||
{"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, "stop_reason": ",5", "token_ids": [20]}
|
||||
{"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras}
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -545,3 +545,31 @@ def test_provider_refusal_string_is_preserved_on_message() -> None:
|
|||
"object": "chat.completion",
|
||||
}
|
||||
assert model_response.choices[0].message.refusal == "I can't help with that."
|
||||
|
||||
|
||||
def test_vllm_null_stop_reason_is_kept_only_on_the_choice_that_carries_it() -> None:
|
||||
model_response: Final = convert_to_model_response_object(
|
||||
response_object={
|
||||
"id": "chatcmpl-vllm-natural-stop",
|
||||
"created": 4,
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"object": "chat.completion",
|
||||
"choices": [
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "1, 2"},
|
||||
"stop_reason": None,
|
||||
},
|
||||
{"finish_reason": "stop", "index": 1, "message": {"role": "assistant", "content": "1, 2"}},
|
||||
],
|
||||
},
|
||||
model_response_object=ModelResponse(),
|
||||
)
|
||||
|
||||
dumped_choices: Final = model_dump_with_preserved_fields(model_response, exclude_unset=True)["choices"]
|
||||
|
||||
assert [(choice["index"], "stop_reason" in choice, choice.get("stop_reason")) for choice in dumped_choices] == [
|
||||
(0, True, None),
|
||||
(1, False, None),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -13615,6 +13615,11 @@ def test_is_injected_stream_usage_artifact():
|
|||
)
|
||||
assert _is_injected_stream_usage_artifact(null_token_ids_chunk) is True
|
||||
|
||||
unrelated_extra_chunk: Final = ModelResponseStream(
|
||||
model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), text="")]
|
||||
)
|
||||
assert _is_injected_stream_usage_artifact(unrelated_extra_chunk) is True
|
||||
|
||||
|
||||
async def _collect_async_data_generator_frames(request_data: dict) -> list:
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from collections.abc import AsyncIterator, Iterator
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -8,7 +9,7 @@ from litellm.exceptions import APIConnectionError
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
|
||||
from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage
|
||||
from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage
|
||||
|
||||
|
||||
def test_reasoning_content_preserved_in_text_completion_wrapper():
|
||||
|
|
@ -98,8 +99,7 @@ def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: obj
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"])
|
||||
def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None:
|
||||
def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]:
|
||||
logging_obj: Final = Logging(
|
||||
model="m",
|
||||
messages=[{"role": "user", "content": "1, 2, 3, 4"}],
|
||||
|
|
@ -109,15 +109,20 @@ def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text
|
|||
litellm_call_id="call-1",
|
||||
function_id="function-1",
|
||||
)
|
||||
upstream: Final = (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7"))
|
||||
chat_stream: Final = CustomStreamWrapper(
|
||||
completion_stream=iter(upstream),
|
||||
model="m",
|
||||
custom_llm_provider="text-completion-openai",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m"))
|
||||
|
||||
responses: Final = list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m"))
|
||||
|
||||
@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"])
|
||||
def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None:
|
||||
responses: Final = _text_completion_stream(
|
||||
(_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7"))
|
||||
)
|
||||
|
||||
assert responses[-1].choices[0].model_dump(exclude_none=True) == {
|
||||
"finish_reason": "stop",
|
||||
|
|
@ -126,6 +131,22 @@ def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("final_text", [", 6", ""], ids=["text-then-finish", "finish-only"])
|
||||
def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(final_text: str) -> None:
|
||||
responses: Final = _text_completion_stream(
|
||||
(
|
||||
_completion_chunk("", None, token_ids=[11]),
|
||||
_completion_chunk("", None, token_ids=[220]),
|
||||
_completion_chunk(", 5", None, token_ids=[20]),
|
||||
_completion_chunk(final_text, "stop", stop_reason=", 7", token_ids=[11, 220, 21]),
|
||||
)
|
||||
)
|
||||
choices: Final = [response.choices[0] for response in responses]
|
||||
|
||||
assert list(chain.from_iterable(choice.get("token_ids") or () for choice in choices)) == [11, 220, 20, 11, 220, 21]
|
||||
assert [choice.get("stop_reason") for choice in choices if choice.finish_reason is not None] == [", 7"]
|
||||
|
||||
|
||||
def test_text_completion_conversion_rejects_non_streaming_choices() -> None:
|
||||
chunk: Final = Choices()
|
||||
wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue