mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(hosted_vllm): keep stop_reason on the last content chunk, keep both fields in built stream responses, and key the cache on return_token_ids
This commit is contained in:
mateo-berri
parent
4b1069df61
commit
31ae498b9b
6 changed files with 144 additions and 11 deletions
|
|
@ -42,7 +42,7 @@ from .redis_semantic_cache import RedisSemanticCache
|
|||
from .s3_cache import S3Cache
|
||||
|
||||
RESPONSE_CACHE_TARGET: Final = "llm_response"
|
||||
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
|
||||
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token", "return_token_ids"})
|
||||
_DECODED_LITERAL: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import VLLM_CHOICE_FIELDS
|
||||
from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
|
||||
|
|
@ -454,6 +455,7 @@ class ChunkProcessor:
|
|||
"index": 0,
|
||||
"message": {"role": role, "content": ""},
|
||||
"finish_reason": finish_reason,
|
||||
**_streamed_vllm_choice_fields(_first_choices_of_index_zero(chunks)),
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
|
|
@ -1306,6 +1308,7 @@ def _rebuilt_text_choice(index: int, choices: Sequence[TextChoices]) -> TextChoi
|
|||
index=index,
|
||||
logprobs=_merged_logprobs(choices),
|
||||
finish_reason=next((choice.finish_reason for choice in reversed(choices) if choice.finish_reason), None),
|
||||
**_streamed_vllm_choice_fields(choices),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1439,6 +1442,32 @@ def _field_of(value: object, name: str) -> object:
|
|||
return getattr(value, name, None)
|
||||
|
||||
|
||||
def _first_choices_of_index_zero(chunks: Iterable[object]) -> tuple[object, ...]:
|
||||
first_choices: Final = (next(iter(object_items(_field_of(chunk, "choices"))), None) for chunk in chunks)
|
||||
return tuple(choice for choice in first_choices if _field_of(choice, "index") in (0, None))
|
||||
|
||||
|
||||
def _choice_extras(choice: object) -> Mapping[str, object]:
|
||||
if isinstance(choice, BaseModel):
|
||||
return choice.model_extra or {}
|
||||
return object_mapping(choice) or {}
|
||||
|
||||
|
||||
def _streamed_vllm_choice_fields(choices: Iterable[object]) -> Mapping[str, object]:
|
||||
extras: Final = tuple(map(_choice_extras, choices))
|
||||
token_id_runs: Final = tuple(object_items(fields.get("token_ids")) for fields in extras)
|
||||
stop_reasons: Final = tuple(reason for fields in extras if (reason := fields.get("stop_reason")) is not None)
|
||||
has_token_ids: Final = any(fields.get("token_ids") is not None for fields in extras)
|
||||
streamed_values: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
"token_ids": list(chain.from_iterable(token_id_runs)) if has_token_ids else None,
|
||||
"stop_reason": stop_reasons[-1] if stop_reasons else None,
|
||||
}
|
||||
)
|
||||
present_fields: Final = frozenset(chain.from_iterable(fields.keys() for fields in extras))
|
||||
return MappingProxyType({field: streamed_values[field] for field in VLLM_CHOICE_FIELDS if field in present_fields})
|
||||
|
||||
|
||||
def _delta_reasoning_items(delta: object) -> Sequence[object]:
|
||||
items: Final = _field_of(delta, "reasoning_items")
|
||||
return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim
|
||||
|
|
|
|||
|
|
@ -1225,14 +1225,14 @@ class CustomStreamWrapper:
|
|||
_prepend_upstream_choice_extras(
|
||||
model_response.choices[0],
|
||||
upstream_choice,
|
||||
upstream_choice_extra_keys - {"message", *STREAM_FINISH_CHOICE_FIELDS},
|
||||
upstream_choice_extra_keys - {"message"},
|
||||
)
|
||||
else:
|
||||
choices: Final = []
|
||||
for choice in upstream_chunk.choices:
|
||||
try:
|
||||
if isinstance(choice, BaseModel):
|
||||
choice_json = choice.model_dump(exclude=set(STREAM_FINISH_CHOICE_FIELDS))
|
||||
choice_json = choice.model_dump()
|
||||
choice_json.pop(
|
||||
"finish_reason", None
|
||||
) # for mistral etc. which return a value in their last chunk (not-openai compatible).
|
||||
|
|
@ -1316,10 +1316,12 @@ class CustomStreamWrapper:
|
|||
model_response.choices[0].finish_reason = map_finish_reason(
|
||||
finish_reason=self.received_finish_reason
|
||||
) # ensure consistent output to openai
|
||||
if _upstream_choice_index(_original_chunk) == model_response.choices[0].index:
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
if response_obj.get("is_finished"):
|
||||
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
if (
|
||||
response_obj.get("is_finished")
|
||||
and _upstream_choice_index(_original_chunk) == model_response.choices[0].index
|
||||
):
|
||||
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
|
||||
|
||||
self.sent_last_chunk = True
|
||||
|
||||
|
|
|
|||
|
|
@ -463,6 +463,17 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token()
|
|||
assert cache.get_cache_key(**request, continuation_token="") == first_segment_key
|
||||
|
||||
|
||||
def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None:
|
||||
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
|
||||
request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]}
|
||||
|
||||
plain_key: Final = cache.get_cache_key(**request)
|
||||
|
||||
assert cache.get_cache_key(**request, return_token_ids=True) != plain_key
|
||||
assert cache.get_cache_key(**request, return_token_ids=False) == plain_key
|
||||
assert cache.get_cache_key(**request, return_token_ids=None) == plain_key
|
||||
|
||||
|
||||
class PhaseRecordingCache(InMemoryCache):
|
||||
"""Records the target and the active span each read / write ran under, as a Redis span would."""
|
||||
|
||||
|
|
|
|||
|
|
@ -2519,3 +2519,73 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it
|
|||
(entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
|
||||
assert entry["guardrail_usage"] == {"topicPolicyUnits": 2}
|
||||
assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015)
|
||||
|
||||
|
||||
def _vllm_chat_chunk(index: int, content: str, finish_reason: str | None = None, **fields: object) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-1",
|
||||
created=1,
|
||||
model="qwen3",
|
||||
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason, **fields)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("finish_content", ["", "4"], ids=["finish-only", "content-then-finish"])
|
||||
def test_stream_chunk_builder_keeps_the_vllm_stop_reason_and_every_token_id_of_the_first_choice(
|
||||
finish_content: str,
|
||||
) -> None:
|
||||
built: Final = stream_chunk_builder(
|
||||
chunks=[
|
||||
_vllm_chat_chunk(0, "1,", stop_reason=None, token_ids=[16, 11]),
|
||||
_vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]),
|
||||
_vllm_chat_chunk(0, "2,3", stop_reason=None, token_ids=[17, 11, 18]),
|
||||
_vllm_chat_chunk(0, finish_content, "stop", stop_reason=",5", token_ids=[19]),
|
||||
]
|
||||
)
|
||||
|
||||
assert built is not None
|
||||
assert built.model_dump()["choices"][0]["stop_reason"] == ",5"
|
||||
assert built.model_dump()["choices"][0]["token_ids"] == [16, 11, 17, 11, 18, 19]
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_null_vllm_fields_and_adds_none_for_other_providers() -> None:
|
||||
vllm_built: Final = stream_chunk_builder(
|
||||
chunks=[
|
||||
_vllm_chat_chunk(0, "hi", stop_reason=None, token_ids=None),
|
||||
_vllm_chat_chunk(0, "", "stop", stop_reason=None, token_ids=None),
|
||||
]
|
||||
)
|
||||
other_built: Final = stream_chunk_builder(
|
||||
chunks=[_vllm_chat_chunk(0, "hi"), _vllm_chat_chunk(0, "", "stop")]
|
||||
)
|
||||
|
||||
assert vllm_built is not None and other_built is not None
|
||||
assert {
|
||||
field: vllm_built.model_dump()["choices"][0][field] for field in ("stop_reason", "token_ids")
|
||||
} == {"stop_reason": None, "token_ids": None}
|
||||
assert {"stop_reason", "token_ids"}.isdisjoint(other_built.model_dump()["choices"][0])
|
||||
|
||||
|
||||
def test_stream_chunk_builder_text_completion_keeps_each_choices_vllm_stop_reason_and_token_ids() -> None:
|
||||
def text_chunk(index: int, text: str, **fields: object) -> TextCompletionResponse:
|
||||
return TextCompletionResponse(
|
||||
id="cmpl-1",
|
||||
object="text_completion",
|
||||
created=1,
|
||||
model="qwen3",
|
||||
choices=[TextChoices(index=index, text=text, **fields)],
|
||||
)
|
||||
|
||||
built: Final = stream_chunk_builder_text_completion(
|
||||
chunks=[
|
||||
text_chunk(0, "1,", token_ids=[16, 11], stop_reason=None),
|
||||
text_chunk(1, "7,", token_ids=[22, 11], stop_reason=None),
|
||||
text_chunk(0, "2", token_ids=[17], stop_reason=", 3"),
|
||||
text_chunk(1, "8", token_ids=[23], stop_reason=151645),
|
||||
]
|
||||
)
|
||||
|
||||
assert [
|
||||
(choice["index"], choice["stop_reason"], choice["token_ids"]) for choice in built.model_dump()["choices"]
|
||||
] == [(0, ", 3", [16, 11, 17]), (1, 151645, [22, 11, 23])]
|
||||
|
||||
|
|
|
|||
|
|
@ -6103,11 +6103,12 @@ def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id(
|
|||
),
|
||||
)
|
||||
|
||||
assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason]
|
||||
assert [
|
||||
(choice["finish_reason"], choice.get("stop_reason"))
|
||||
(_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"])
|
||||
for choice in emitted
|
||||
if choice["finish_reason"] is not None or choice.get("stop_reason") is not None
|
||||
] == [("stop", stop_reason)]
|
||||
if choice["finish_reason"] is None and choice.get("stop_reason") is not None
|
||||
] == ([(final_content, stop_reason)] if final_content else [])
|
||||
assert _joined_token_ids(emitted) == (16, 11, 17, 20)
|
||||
|
||||
|
||||
|
|
@ -6217,6 +6218,27 @@ def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_
|
|||
] == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("first_finished", [0, 1])
|
||||
@pytest.mark.parametrize(
|
||||
"finish_contents", [("2", ""), ("2", "2")], ids=["content-then-finish-only", "both-content-then-finish"]
|
||||
)
|
||||
def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final_content_chunk(
|
||||
logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str]
|
||||
) -> None:
|
||||
emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents))
|
||||
content_chunk_stop_reasons: Final = frozenset(
|
||||
(choice["index"], choice["stop_reason"])
|
||||
for choice in emitted
|
||||
if choice["finish_reason"] is None and choice.get("stop_reason") is not None
|
||||
)
|
||||
finished_with_content: Final = frozenset(
|
||||
(index, _TWO_CHOICE_STOP_REASONS[index])
|
||||
for index, content in zip((first_finished, 1 - first_finished), finish_contents)
|
||||
if content
|
||||
)
|
||||
|
||||
assert content_chunk_stop_reasons == finished_with_content
|
||||
|
||||
_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19)
|
||||
|
||||
|
||||
|
|
@ -6264,7 +6286,6 @@ async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_i
|
|||
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 if choice["finish_reason"] is None} == {None}
|
||||
assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue