mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
Merge a7ac7ad54e into 57c1df8c33
This commit is contained in:
Prabal Pratap Singh (ZeroX)
GitHub
commit
b48b6281bf
10 changed files with 334 additions and 8 deletions
|
|
@ -686,6 +686,8 @@ STREAM_SSE_DONE_STRING: Final[str] = "[DONE]"
|
|||
STREAM_SSE_DATA_PREFIX: Final[str] = "data: "
|
||||
STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n'
|
||||
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",)
|
||||
### SPEND TRACKING ###
|
||||
DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float(
|
||||
"DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict
|
|||
import litellm
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_FIELDS
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost
|
||||
from litellm.litellm_core_utils.hidden_params import (
|
||||
|
|
@ -113,6 +114,41 @@ def _prepend_upstream_choice_extras(
|
|||
choice.__pydantic_fields_set__.update(upstream_extras)
|
||||
|
||||
|
||||
_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None:
|
||||
match getattr(upstream_chunk, "choices", None):
|
||||
case [BaseModel() as first_choice, *_]:
|
||||
return first_choice
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
|
||||
first_choice: Final = _upstream_first_choice(upstream_chunk)
|
||||
if first_choice is None:
|
||||
return _NO_UPSTREAM_CHOICE_FIELDS
|
||||
upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {}
|
||||
return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras})
|
||||
|
||||
|
||||
def _upstream_choice_index(upstream_chunk: object) -> int:
|
||||
upstream_index: Final[object] = getattr(_upstream_first_choice(upstream_chunk), "index", 0)
|
||||
return upstream_index if isinstance(upstream_index, int) else 0
|
||||
|
||||
|
||||
def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{field: value for field, value in _upstream_choice_fields(upstream_chunk, fields).items() if value is not None}
|
||||
)
|
||||
|
||||
|
||||
def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) -> None:
|
||||
for field, value in fields.items():
|
||||
setattr(choice, field, value)
|
||||
|
||||
|
||||
def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool:
|
||||
return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields)
|
||||
|
||||
|
|
@ -337,6 +373,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.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
|
||||
self.special_tokens = [
|
||||
"<|assistant|>",
|
||||
|
|
@ -1268,12 +1305,26 @@ 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)
|
||||
if response_obj.get("is_finished"):
|
||||
_apply_choice_fields(
|
||||
model_response.choices[0],
|
||||
_non_null_upstream_choice_fields(_original_chunk, STREAM_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
|
||||
):
|
||||
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)
|
||||
return model_response
|
||||
else:
|
||||
if hasattr(model_response, "usage"):
|
||||
self.chunks.append(model_response)
|
||||
|
|
@ -1497,6 +1548,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)
|
||||
if response_obj["usage"] is not None:
|
||||
_text_completion_usage: Final[Usage] = response_obj["usage"]
|
||||
setattr(
|
||||
|
|
@ -1586,6 +1638,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)
|
||||
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)
|
||||
|
|
@ -1931,6 +1984,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)
|
||||
|
||||
## if tool use
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -3,8 +3,11 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions`
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Coroutine
|
||||
from typing import Final, Literal, cast, overload
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_image_mime_type_from_url,
|
||||
|
|
@ -22,10 +25,24 @@ 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]]:
|
||||
|
|
@ -106,6 +123,38 @@ 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]:
|
||||
|
|
|
|||
|
|
@ -10065,6 +10065,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 _fast_serialize_simple_model_response_stream(
|
||||
chunk: ModelResponseStream,
|
||||
) -> bytes | None:
|
||||
|
|
@ -10087,6 +10090,8 @@ def _fast_serialize_simple_model_response_stream(
|
|||
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:
|
||||
return None
|
||||
|
||||
delta: Final[object] = getattr(choice, "delta", None)
|
||||
if delta is None:
|
||||
|
|
@ -10165,6 +10170,11 @@ 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):
|
||||
return False
|
||||
delta: Final = choice.delta if hasattr(choice, "delta") else None
|
||||
if delta is None:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.constants import STREAM_FINISH_CHOICE_FIELDS
|
||||
from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse
|
||||
|
||||
|
||||
|
|
@ -34,11 +35,15 @@ class TextCompletionStreamWrapper:
|
|||
raise Exception
|
||||
|
||||
if chunk["choices"]:
|
||||
delta: Final = chunk["choices"][0]["delta"]
|
||||
chat_choice: Final = chunk["choices"][0]
|
||||
delta: Final = chat_choice["delta"]
|
||||
text_choices["text"] = delta["content"]
|
||||
text_choices["reasoning_content"] = delta.get("reasoning_content")
|
||||
text_choices["index"] = chunk["choices"][0]["index"]
|
||||
text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"]
|
||||
text_choices["index"] = chat_choice["index"]
|
||||
text_choices["finish_reason"] = chat_choice["finish_reason"]
|
||||
for field in STREAM_FINISH_CHOICE_FIELDS:
|
||||
if field in chat_choice:
|
||||
text_choices[field] = chat_choice[field]
|
||||
response["choices"] = [text_choices]
|
||||
else:
|
||||
response["choices"] = []
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import threading
|
|||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from types import CodeType, FrameType, MappingProxyType
|
||||
from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
_ProviderChunkParsed,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
|
||||
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -6158,6 +6160,89 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj:
|
|||
assert result.choices[0].delta.content == "hello"
|
||||
|
||||
|
||||
def _hosted_vllm_stream_choices(
|
||||
logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]]
|
||||
) -> tuple[dict[str, object], ...]:
|
||||
sse_lines: Final = (
|
||||
*(
|
||||
"data: " + json.dumps({"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [c]})
|
||||
for c in upstream_choices
|
||||
),
|
||||
"data: [DONE]",
|
||||
)
|
||||
completion_stream: Final = HostedVLLMChatConfig().get_model_response_iterator(
|
||||
streaming_response=iter(sse_lines), sync_stream=True, json_mode=False
|
||||
)
|
||||
return tuple(
|
||||
chunk.choices[0].model_dump()
|
||||
for chunk in CustomStreamWrapper(
|
||||
completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stop_reason", [",5", 151645])
|
||||
@pytest.mark.parametrize("final_content", [",3", ""], ids=["content-then-finish", "finish-only"])
|
||||
def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id(
|
||||
logging_obj: Logging, final_content: str, stop_reason: str | int
|
||||
) -> None:
|
||||
emitted: Final = _hosted_vllm_stream_choices(
|
||||
logging_obj,
|
||||
(
|
||||
{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},
|
||||
{"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]},
|
||||
{"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]},
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": final_content},
|
||||
"finish_reason": "stop",
|
||||
"stop_reason": stop_reason,
|
||||
"token_ids": [20],
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
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]
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(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": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]},
|
||||
{"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [17]},
|
||||
{"index": 1, "delta": {"content": "2"}, "finish_reason": None, "token_ids": [18]},
|
||||
),
|
||||
)
|
||||
|
||||
assert [(choice["index"], choice.get("token_ids")) for choice in emitted if choice.get("token_ids")] == [
|
||||
(1, [16]),
|
||||
(0, [17]),
|
||||
(1, [18]),
|
||||
]
|
||||
|
||||
|
||||
def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]:
|
||||
return (
|
||||
{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None},
|
||||
{"index": 0, "delta": {"content": ""}, "finish_reason": None, **token_ids},
|
||||
{"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, **token_ids},
|
||||
{"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "stop_reason": ",3", **token_ids},
|
||||
)
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without_them(logging_obj: Logging) -> None:
|
||||
with_null_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({"token_ids": None}))
|
||||
without_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({}))
|
||||
|
||||
assert [(c["delta"], c["finish_reason"]) for c in with_null_ids] == [
|
||||
(c["delta"], c["finish_reason"]) for c in without_ids
|
||||
]
|
||||
|
||||
|
||||
def _later_chunk_outputs(
|
||||
custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging
|
||||
) -> tuple[ModelResponseStream, ...]:
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
|
||||
from litellm.constants import (
|
||||
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
|
||||
)
|
||||
from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
def test_hosted_vllm_chat_transformation_file_url():
|
||||
|
|
@ -397,3 +399,61 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema():
|
|||
assert tools[0]["function"]["name"] == "search"
|
||||
assert tools[0]["function"]["description"] == "Search docs"
|
||||
assert tools[0]["function"]["parameters"] == input_schema
|
||||
|
||||
|
||||
def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse:
|
||||
payload: Final = {
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "glm",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "1, 2, 3, 4"},
|
||||
"logprobs": None,
|
||||
"finish_reason": "stop",
|
||||
**choice_extras,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
|
||||
}
|
||||
raw_response: Final = httpx.Response(
|
||||
200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions")
|
||||
)
|
||||
return HostedVLLMChatConfig().transform_response(
|
||||
model="hosted_vllm/glm",
|
||||
raw_response=raw_response,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"choice_extras",
|
||||
[
|
||||
{"stop_reason": ", 5", "token_ids": [16, 11, 220]},
|
||||
{"stop_reason": 128009},
|
||||
{"stop_reason": None},
|
||||
],
|
||||
)
|
||||
def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]):
|
||||
response: Final = _transform_vllm_chat_response(choice_extras)
|
||||
|
||||
dumped_choice: Final[dict[str, dict[str, object]]] = json.loads(response.model_dump_json())["choices"][0]
|
||||
for field, value in choice_extras.items():
|
||||
assert dumped_choice[field] == value
|
||||
assert dumped_choice["provider_specific_fields"][field] == value
|
||||
|
||||
|
||||
def test_hosted_vllm_does_not_invent_vllm_choice_fields():
|
||||
response: Final = _transform_vllm_chat_response({})
|
||||
|
||||
dumped_choice: Final[dict[str, object]] = json.loads(response.model_dump_json())["choices"][0]
|
||||
assert "stop_reason" not in dumped_choice
|
||||
assert "token_ids" not in dumped_choice
|
||||
|
|
|
|||
|
|
@ -576,6 +576,16 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes():
|
|||
}
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_keeps_provider_choice_fields():
|
||||
chunk: Final = _simple_chunk(content="1,2,3,4")
|
||||
chunk.choices[0]["stop_reason"] = ",5"
|
||||
chunk.choices[0]["token_ids"] = [20]
|
||||
|
||||
assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [
|
||||
{"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, "stop_reason": ",5", "token_ids": [20]}
|
||||
]
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_invalid_input_raises_attribute_error():
|
||||
"""The helper is typed as ``BaseModel`` — handing it a plain dict trips
|
||||
the attribute-access path (no ``model_dump_json``)."""
|
||||
|
|
|
|||
|
|
@ -13582,7 +13582,7 @@ def _stream_usage_generator_chunks():
|
|||
|
||||
def test_is_injected_stream_usage_artifact():
|
||||
from litellm.proxy.proxy_server import _is_injected_stream_usage_artifact
|
||||
from litellm.types.utils import ModelResponseStream, Usage
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
|
||||
|
||||
content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks()
|
||||
assert _is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True
|
||||
|
|
@ -13605,6 +13605,16 @@ def test_is_injected_stream_usage_artifact():
|
|||
|
||||
assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False
|
||||
|
||||
held_back_token_ids_chunk: Final = ModelResponseStream(
|
||||
model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])]
|
||||
)
|
||||
assert _is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False
|
||||
|
||||
null_token_ids_chunk: Final = ModelResponseStream(
|
||||
model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)]
|
||||
)
|
||||
assert _is_injected_stream_usage_artifact(null_token_ids_chunk) is True
|
||||
|
||||
|
||||
async def _collect_async_data_generator_frames(request_data: dict) -> list:
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
|
|
|
|||
|
|
@ -2,8 +2,11 @@ from collections.abc import AsyncIterator, Iterator
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from openai.types.completion import Completion
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -85,6 +88,44 @@ async def test_text_completion_stream_usage_is_opt_in(
|
|||
]
|
||||
|
||||
|
||||
def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: object) -> Completion:
|
||||
return Completion.construct(
|
||||
id="cmpl-1",
|
||||
created=1,
|
||||
model="m",
|
||||
object="text_completion",
|
||||
choices=[{"index": 0, "text": text, "finish_reason": finish_reason, "logprobs": None, **choice_extras}],
|
||||
)
|
||||
|
||||
|
||||
@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:
|
||||
logging_obj: Final = Logging(
|
||||
model="m",
|
||||
messages=[{"role": "user", "content": "1, 2, 3, 4"}],
|
||||
stream=True,
|
||||
call_type="text_completion",
|
||||
start_time=0.0,
|
||||
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,
|
||||
)
|
||||
|
||||
responses: Final = list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m"))
|
||||
|
||||
assert responses[-1].choices[0].model_dump(exclude_none=True) == {
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"stop_reason": ", 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