mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix: Preserve deferred streaming choice metadata
This commit is contained in:
Alberto Perdomo
parent
6a545a1a52
commit
153342dca8
4 changed files with 108 additions and 18 deletions
|
|
@ -56,10 +56,10 @@ from litellm.types.utils import Logprobs as TextCompletionLogprobs
|
|||
from .get_headers import get_response_headers
|
||||
|
||||
_MESSAGE_FIELDS: Final[frozenset] = frozenset(Message.model_fields.keys())
|
||||
_CHOICES_FIELDS: Final[frozenset] = frozenset(Choices.model_fields.keys())
|
||||
_CHOICES_EXPLICIT_KWARGS: Final[frozenset] = _CHOICES_FIELDS | {"enhancements"}
|
||||
_CHOICES_FIELDS: Final[frozenset[str]] = frozenset(Choices.model_fields.keys())
|
||||
_CHOICES_EXPLICIT_KWARGS: Final[frozenset[str]] = _CHOICES_FIELDS | {"enhancements"}
|
||||
_MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"}
|
||||
_CACHED_CHOICE_RESERVED_KEYS: Final[frozenset] = STREAMING_CHOICE_OPENAI_FIELDS | {"message", "finish_details"}
|
||||
_CACHED_CHOICE_RESERVED_KEYS: Final[frozenset[str]] = STREAMING_CHOICE_OPENAI_FIELDS | {"message", "finish_details"}
|
||||
|
||||
|
||||
def _normalize_images_for_message(
|
||||
|
|
@ -148,8 +148,8 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]:
|
|||
return _REPLAY_CONTENT_SLICE_RE.findall(content)
|
||||
|
||||
|
||||
def _cached_choice_provider_extras(choice: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {key: value for key, value in choice.items() if key not in _CACHED_CHOICE_RESERVED_KEYS}
|
||||
def _choice_provider_extras(choice: Mapping[str, object], reserved_keys: frozenset[str]) -> Mapping[str, object]:
|
||||
return {key: value for key, value in choice.items() if key not in reserved_keys}
|
||||
|
||||
|
||||
def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
|
||||
|
|
@ -260,7 +260,7 @@ async def convert_to_streaming_response_async(
|
|||
index=idx,
|
||||
delta=delta,
|
||||
logprobs=logprobs,
|
||||
**_cached_choice_provider_extras(choice),
|
||||
**_choice_provider_extras(choice, _CACHED_CHOICE_RESERVED_KEYS),
|
||||
)
|
||||
choice_list.append(choice)
|
||||
|
||||
|
|
@ -360,7 +360,7 @@ def convert_to_streaming_response(
|
|||
delta=delta,
|
||||
logprobs=logprobs,
|
||||
enhancements=enhancements,
|
||||
**_cached_choice_provider_extras(choice),
|
||||
**_choice_provider_extras(choice, _CACHED_CHOICE_RESERVED_KEYS),
|
||||
)
|
||||
|
||||
choice_list.append(choice)
|
||||
|
|
@ -814,7 +814,7 @@ def convert_to_model_response_object(
|
|||
logprobs=logprobs,
|
||||
enhancements=enhancements,
|
||||
provider_specific_fields=provider_specific_fields,
|
||||
**{f: v for f, v in provider_specific_fields.items() if f not in _CHOICES_EXPLICIT_KWARGS},
|
||||
**_choice_provider_extras(provider_specific_fields, _CHOICES_EXPLICIT_KWARGS),
|
||||
)
|
||||
if provider_sent_logprobs and logprobs is None:
|
||||
choice.logprobs = None
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.types.llms.base import LiteLLMBaseModel
|
|||
from litellm.types.llms.openai import OpenAIChatCompletionChunk
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
STREAMING_CHOICE_OPENAI_FIELDS,
|
||||
CacheCreationTokenDetails,
|
||||
CompletionTokensDetailsWrapper,
|
||||
Delta,
|
||||
|
|
@ -257,6 +258,7 @@ class CustomStreamWrapper:
|
|||
self.system_fingerprint: str | None = None
|
||||
self._provider_response_model: str | None = None
|
||||
self.received_finish_reason: str | None = None
|
||||
self._deferred_finish_chunk: ModelResponseStream | None = None
|
||||
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
|
||||
self.special_tokens = [
|
||||
"<|assistant|>",
|
||||
|
|
@ -293,7 +295,9 @@ class CustomStreamWrapper:
|
|||
self.received_model_id: str | None = None
|
||||
self.logging_loop = None
|
||||
self.rules = Rules()
|
||||
self.stream_options = stream_options or getattr(logging_obj, "stream_options", None)
|
||||
self.stream_options: Mapping[str, object] | None = stream_options or getattr(
|
||||
logging_obj, "stream_options", None
|
||||
)
|
||||
self.messages = getattr(logging_obj, "messages", None)
|
||||
self.sent_stream_usage = False
|
||||
self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
|
||||
|
|
@ -301,7 +305,7 @@ class CustomStreamWrapper:
|
|||
isinstance(self.stream_options, dict) and self.stream_options.get("continuous_usage_stats") is True
|
||||
) and self.send_stream_usage
|
||||
self.tool_call = False
|
||||
self.chunks: list = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options
|
||||
self.chunks: list[ModelResponseStream] = []
|
||||
self._repeated_messages_count = 1
|
||||
self.is_function_call = self.check_is_function_call(logging_obj=logging_obj)
|
||||
self.created: int | None = None
|
||||
|
|
@ -461,7 +465,7 @@ class CustomStreamWrapper:
|
|||
)
|
||||
self._restore_consumer_correlation_context()
|
||||
|
||||
def check_send_stream_usage(self, stream_options: dict | None):
|
||||
def check_send_stream_usage(self, stream_options: Mapping[str, object] | None) -> bool:
|
||||
return stream_options is not None and stream_options.get("include_usage", False) is True
|
||||
|
||||
def check_is_function_call(self, logging_obj) -> bool:
|
||||
|
|
@ -509,9 +513,7 @@ class CustomStreamWrapper:
|
|||
|
||||
last_content: Final = self.chunks[-1].choices[0].delta.content
|
||||
|
||||
if (
|
||||
last_content is None or not isinstance(last_content, str) or len(last_content) <= 2
|
||||
): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946
|
||||
if last_content is None or len(last_content) <= 2:
|
||||
self._repeated_messages_count = 1
|
||||
return
|
||||
|
||||
|
|
@ -1070,6 +1072,24 @@ class CustomStreamWrapper:
|
|||
|
||||
self._optional_combine_thinking_block_in_choices(model_response=model_response)
|
||||
|
||||
if self.received_finish_reason is not None:
|
||||
self._deferred_finish_chunk = ModelResponseStream(choices=model_response.choices).model_copy(
|
||||
deep=True
|
||||
)
|
||||
return model_response.model_copy(
|
||||
update={
|
||||
"choices": [
|
||||
StreamingChoices(
|
||||
**{
|
||||
key: value
|
||||
for key, value in choice.model_dump().items()
|
||||
if key in STREAMING_CHOICE_OPENAI_FIELDS
|
||||
}
|
||||
)
|
||||
for choice in model_response.choices
|
||||
]
|
||||
}
|
||||
)
|
||||
return model_response
|
||||
else:
|
||||
return
|
||||
|
|
@ -1773,9 +1793,13 @@ class CustomStreamWrapper:
|
|||
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
|
||||
|
||||
def finish_reason_handler(self):
|
||||
from litellm.litellm_core_utils.core_helpers import preserve_streaming_choice_provider_extras
|
||||
|
||||
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}
|
||||
)
|
||||
if self._deferred_finish_chunk is not None:
|
||||
preserve_streaming_choice_provider_extras(model_response, self._deferred_finish_chunk)
|
||||
_finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason
|
||||
if _finish_reason is not None:
|
||||
model_response.choices[0].finish_reason = _finish_reason
|
||||
|
|
@ -2559,7 +2583,7 @@ def _coerce_token_details(
|
|||
return details_type(**(raw if isinstance(raw, dict) else raw.model_dump()))
|
||||
|
||||
|
||||
def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
|
||||
def calculate_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage:
|
||||
"""Assume most recent usage chunk has total usage uptil then."""
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
|
||||
attach_cache_creation_token_details,
|
||||
|
|
|
|||
|
|
@ -9858,7 +9858,7 @@ def _restamp_streaming_chunk_model(
|
|||
|
||||
|
||||
def _streaming_choice_carries_provider_extras(choice: object) -> bool:
|
||||
extra: Final = getattr(choice, "__pydantic_extra__", None)
|
||||
extra: Final[Mapping[str, object] | None] = getattr(choice, "__pydantic_extra__", None)
|
||||
if not isinstance(extra, dict):
|
||||
return False
|
||||
return any(key not in STREAMING_CHOICE_OPENAI_FIELDS for key in extra)
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@ import asyncio
|
|||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, NoReturn, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -6691,11 +6693,15 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream():
|
|||
assert stream.close_calls == 1
|
||||
|
||||
|
||||
def _vllm_wrapper(chunks: list[ModelResponseStream], stream_options: dict | None = None) -> CustomStreamWrapper:
|
||||
def _vllm_wrapper(
|
||||
chunks: Sequence[ModelResponseStream | ChatCompletionChunk],
|
||||
stream_options: dict | None = None,
|
||||
custom_llm_provider: str = "hosted_vllm",
|
||||
) -> CustomStreamWrapper:
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=ModelResponseListIterator(model_responses=chunks),
|
||||
model="hosted_vllm/m",
|
||||
custom_llm_provider="hosted_vllm",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream_options=stream_options,
|
||||
logging_obj=Logging(
|
||||
model="hosted_vllm/m",
|
||||
|
|
@ -6732,6 +6738,66 @@ async def test_vllm_choice_extras_survive_the_terminal_finish_chunk(sync_mode: b
|
|||
assert getattr(finish_choice, "token_ids", None) == [9]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.parametrize("source", ["sdk", "http_handler", "cache"])
|
||||
@pytest.mark.parametrize("content", ["", "last token"])
|
||||
@pytest.mark.parametrize(("stop_reason", "token_ids"), [("eos", (9,)), (0, ())])
|
||||
@pytest.mark.asyncio
|
||||
async def test_vllm_finish_choice_extras_survive_serialized_live_and_cached_streams(
|
||||
sync_mode: bool, source: str, content: str, stop_reason: str | int, token_ids: tuple[int, ...]
|
||||
) -> None:
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
convert_to_streaming_response,
|
||||
convert_to_streaming_response_async,
|
||||
)
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
|
||||
from litellm.proxy.proxy_server import _serialize_streaming_chunk
|
||||
|
||||
terminal_choice: Final = {
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"stop_reason": stop_reason,
|
||||
"token_ids": list(token_ids),
|
||||
}
|
||||
cached_response: Final = {
|
||||
"id": "c",
|
||||
"created": 1,
|
||||
"model": "hosted_vllm/m",
|
||||
"choices": [{**terminal_choice, "message": {"role": "assistant", "content": content}}],
|
||||
}
|
||||
replay_chunks: Final = (
|
||||
list(convert_to_streaming_response(cached_response))
|
||||
if sync_mode
|
||||
else [chunk async for chunk in convert_to_streaming_response_async(cached_response)]
|
||||
)
|
||||
parser: Final = OpenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True)
|
||||
live_body: Final = {
|
||||
"id": "c",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "hosted_vllm/m",
|
||||
"choices": [{**terminal_choice, "delta": {"content": content}}],
|
||||
}
|
||||
live_chunk: Final = ChatCompletionChunk(**live_body) if source == "sdk" else parser.chunk_parser(live_body)
|
||||
wrapper: Final = _vllm_wrapper(
|
||||
replay_chunks if source == "cache" else [live_chunk],
|
||||
custom_llm_provider="cached_response" if source == "cache" else "hosted_vllm",
|
||||
)
|
||||
|
||||
emitted: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper]
|
||||
serialized: Final = [json.loads(_serialize_streaming_chunk(chunk)) for chunk in emitted]
|
||||
finish_choices: Final = [chunk["choices"][0] for chunk in serialized if chunk["choices"][0].get("finish_reason")]
|
||||
|
||||
assert len(finish_choices) == 1
|
||||
assert finish_choices[0]["stop_reason"] == terminal_choice["stop_reason"]
|
||||
assert finish_choices[0]["token_ids"] == terminal_choice["token_ids"]
|
||||
assert not finish_choices[0]["delta"].get("content")
|
||||
assert [
|
||||
chunk["choices"][0]["token_ids"] for chunk in serialized if "token_ids" in chunk["choices"][0]
|
||||
] == [terminal_choice["token_ids"]]
|
||||
assert "".join(chunk["choices"][0]["delta"].get("content") or "" for chunk in serialized) == content
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_continuous_usage_stats_keeps_usage_on_every_chunk(sync_mode: bool):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue