mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(hosted_vllm)!: keep vLLM's stop_reason and token_ids on chat choices and stream chunks (#48393)
* feat(hosted_vllm): keep vLLM's stop_reason and token_ids on chat choices Non-streamed hosted_vllm chat completions moved vLLM's choice-level stop_reason and token_ids into provider_specific_fields, so clients written against vLLM could not find them where vLLM puts them. They are now also set on the choice itself, and the copy in provider_specific_fields stays for existing readers. Fixes #47945 * feat(hosted_vllm): keep stop_reason and token_ids on streamed chunks and proxy stream output Streamed chat finish chunks and streamed /v1/completions finish chunks now carry vLLM's stop_reason. Each stream chunk keeps the token_ids vLLM sent with it, including held-back empty chunks, and the proxy no longer strips those chunks as empty usage chunks or drops the fields in its fast serializer. * fix(hosted_vllm): keep the upstream choice index on token-id-only stream chunks * 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. * fix(hosted_vllm): send empty text on token-id chunks, keep fields on their own choice, and keep null stop_reason on slow-path streams Token-id-only chunks now carry an empty content or text string, a finish chunk only takes stop_reason and token_ids from its own upstream choice, content chunks no longer repeat stop_reason, the slow stream serializer keeps a null stop_reason or token_ids, and fields cached under provider_specific_fields come back on the choice * 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 * fix(hosted_vllm): keep another choice's final token_ids and stop_reason on their own stream chunk With n above 1, a finish-only vLLM chunk for a choice other than the one the single finish chunk carries was dropped or taken as choice 0's finish, losing that choice's last token_ids. It now goes out as an id-only chunk on its own choice index. * fix(hosted_vllm): send the first choice's finish chunk before another choice's later token-id chunks * fix(streaming): keep each streamed text completion choice on its own index * fix(streaming): finish each streamed text completion choice on its own index and null token_ids in redacted logs * fix(proxy): drop token_ids from guarded streams, redact choice-level provider fields, and keep one message_delta per Anthropic stream * fix(proxy): skip token-id-only chunks on guarded streams and put stop_reason on the final text chunk * fix(proxy): drop token_ids from the logged and cached response of a guarded stream * fix(proxy): drop token_ids under per-chunk stream checks and redact text-completion stream choices * fix(caching): key return_token_ids and continuation_token sent in extra_body * fix(streaming): drop token ids from streamed chunks while post_call_rules run vLLM holds text back behind a possible stop string but sends its token ids at once, so an SDK stream sent ids before litellm.post_call_rules checked their text. While those rules are set, streamed choices drop token_ids and chunks that carry only ids are skipped, as the proxy already does ahead of a per-chunk check * fix(streaming): keep vLLM fields when a stream's only choice index is not 0 --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Prabal Pratap Singh (ZeroX)
GitHub
parent
61da6fcb4d
commit
b3c397e964
24 changed files with 1656 additions and 39 deletions
|
|
@ -45,13 +45,27 @@ 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)
|
||||
_REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter(
|
||||
list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
|
||||
def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str:
|
||||
raw_extra_body: Final = kwargs.get("extra_body")
|
||||
if not isinstance(raw_extra_body, dict):
|
||||
return ""
|
||||
extra_body: Final = cast( # cast-ok: isinstance leaves the key and value types unknown
|
||||
Mapping[str, object], raw_extra_body
|
||||
)
|
||||
return "".join(
|
||||
f"extra_body.{param}: {extra_body[param]}"
|
||||
for param in sorted(_ALWAYS_KEYED_PROVIDER_PARAMS)
|
||||
if extra_body.get(param)
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def response_cache_phase(operation: Literal["get", "set"]) -> Generator[None]:
|
||||
"""The ``cache.get llm_response`` / ``cache.set llm_response`` span a response-cache read or write runs
|
||||
|
|
@ -429,7 +443,10 @@ class Cache:
|
|||
"" if embeds_thought_signature_in_tool_call_id(kwargs) else "thought_signature_in_tool_call_id: False"
|
||||
)
|
||||
hashed_cache_key = Cache._get_hashed_cache_key(
|
||||
cache_key + reasoning_summary_material + tool_call_id_shape_material
|
||||
cache_key
|
||||
+ _keyed_provider_params_in_extra_body(kwargs)
|
||||
+ reasoning_summary_material
|
||||
+ tool_call_id_shape_material
|
||||
)
|
||||
hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
|
||||
verbose_logger.debug(
|
||||
|
|
|
|||
|
|
@ -699,6 +699,9 @@ 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",)
|
||||
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
|
||||
|
|
|
|||
|
|
@ -918,6 +918,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# post_call guardrails have run; the @client decorator then stores the
|
||||
# enqueue closure here instead of firing it immediately.
|
||||
self.defer_async_logging: bool = False
|
||||
self.drops_stream_token_ids: bool = False
|
||||
self.enqueue_deferred_logging: Callable[[], None] | None = None
|
||||
self._async_success_scheduled: bool = False
|
||||
self._async_success_enqueued: bool = False
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ import json
|
|||
import re
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Generator, Iterator, Mapping, Sequence
|
||||
from collections.abc import Generator, Iterable, Iterator, 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,
|
||||
|
|
@ -170,6 +176,28 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]:
|
|||
return _REPLAY_CONTENT_SLICE_RE.findall(content)
|
||||
|
||||
|
||||
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]:
|
||||
cached_fields: Final = _choice_provider_specific_fields(choice)
|
||||
return MappingProxyType(
|
||||
{
|
||||
field: choice[field] if field in choice else cached_fields[field]
|
||||
for field in VLLM_CHOICE_FIELDS
|
||||
if field in choice or field in cached_fields
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _choice_provider_specific_fields(choice: Mapping[str, object]) -> Mapping[str, object]:
|
||||
fields: Final = choice.get("provider_specific_fields")
|
||||
return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else MappingProxyType({})
|
||||
|
||||
|
||||
_REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"})
|
||||
_REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"}
|
||||
_REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"})
|
||||
|
|
@ -220,8 +248,9 @@ def _replay_frame(
|
|||
choice.finish_reason = finish_reason
|
||||
if not keeps_choice_metadata:
|
||||
choice.logprobs = None
|
||||
if hasattr(choice, "enhancements"):
|
||||
del choice.enhancements
|
||||
_drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS))
|
||||
if finish_reason is None:
|
||||
_drop_choice_fields(choice, STREAM_FINISH_CHOICE_FIELDS)
|
||||
return frame
|
||||
|
||||
|
||||
|
|
@ -325,7 +354,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
|
||||
|
|
@ -388,6 +419,7 @@ def convert_to_streaming_response(
|
|||
delta=delta,
|
||||
logprobs=logprobs,
|
||||
enhancements=enhancements,
|
||||
**_vllm_choice_fields(choice),
|
||||
)
|
||||
|
||||
choice_list.append(choice)
|
||||
|
|
@ -872,6 +904,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
|
||||
|
|
|
|||
|
|
@ -196,6 +196,9 @@ def _redact_function_call(function_call) -> None:
|
|||
|
||||
def _redact_choice_content(choice):
|
||||
"""Helper to redact content in a choice (message or delta)."""
|
||||
if isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices, litellm.utils.TextChoices)):
|
||||
drop_token_ids(choice)
|
||||
_redact_provider_specific_fields(choice)
|
||||
if isinstance(choice, litellm.Choices):
|
||||
if choice.message.content is not None:
|
||||
choice.message.content = REDACTED_BY_LITELLM
|
||||
|
|
@ -216,6 +219,9 @@ def _redact_choice_content(choice):
|
|||
_redact_provider_specific_fields(choice.delta)
|
||||
_redact_tool_calls(getattr(choice.delta, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.delta, "function_call", None))
|
||||
elif isinstance(choice, litellm.utils.TextChoices):
|
||||
_redact_field_when_present(choice, "text")
|
||||
_redact_field_when_present(choice, "reasoning_content")
|
||||
|
||||
|
||||
def _as_mapping(value: object) -> Mapping[str, object] | None:
|
||||
|
|
@ -227,14 +233,27 @@ def _redact_provider_specific_fields(message: object) -> None:
|
|||
setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object
|
||||
|
||||
|
||||
def _drop_provider_specific_token_ids(holder: object) -> None:
|
||||
fields: Final = _as_mapping(getattr(holder, "provider_specific_fields", None))
|
||||
if fields is not None and fields.get("token_ids") is not None:
|
||||
setattr(holder, "provider_specific_fields", {**fields, "token_ids": None}) # noqa: B010 # typed as object
|
||||
|
||||
|
||||
def drop_token_ids(choice: object) -> None:
|
||||
if getattr(choice, "token_ids", None) is not None:
|
||||
setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field
|
||||
for holder in (choice, getattr(choice, "delta", None), getattr(choice, "message", None)):
|
||||
_drop_provider_specific_token_ids(holder)
|
||||
|
||||
|
||||
def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None:
|
||||
for message in (choice.get("message"), choice.get("delta")):
|
||||
for holder in (choice, choice.get("message"), choice.get("delta")):
|
||||
if (
|
||||
(readable := _as_mapping(message)) is not None
|
||||
(readable := _as_mapping(holder)) is not None
|
||||
and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None
|
||||
and isinstance(message, dict)
|
||||
and isinstance(holder, dict)
|
||||
):
|
||||
message["provider_specific_fields"] = _redacted_provider_specific_fields(fields)
|
||||
holder["provider_specific_fields"] = _redacted_provider_specific_fields(fields)
|
||||
|
||||
|
||||
def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
@ -251,6 +270,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[
|
|||
"compaction_blocks",
|
||||
"compaction_start",
|
||||
"compaction_delta",
|
||||
"token_ids",
|
||||
)
|
||||
if key in fields
|
||||
},
|
||||
|
|
@ -422,6 +442,8 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
|
|||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
_redact_provider_specific_fields_dict(choice)
|
||||
if choice.get("token_ids") is not None:
|
||||
choice["token_ids"] = None
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
if choice["message"].get("content") is not None:
|
||||
choice["message"]["content"] = redacted_str
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import VLLM_CHOICE_FIELDS
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices
|
||||
|
||||
|
||||
|
|
@ -9,12 +10,19 @@ def is_empty_streaming_choice(choice: StreamingChoices) -> bool:
|
|||
return False
|
||||
if getattr(choice, "logprobs", None) is not None:
|
||||
return False
|
||||
if not _all_values_none(_vllm_choice_values(choice)):
|
||||
return False
|
||||
delta: Final = choice.delta if hasattr(choice, "delta") else None
|
||||
if delta is None:
|
||||
return True
|
||||
return _all_values_none(delta.model_dump().values())
|
||||
|
||||
|
||||
def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]:
|
||||
choice_extras: Final = choice.model_extra or {}
|
||||
return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS)
|
||||
|
||||
|
||||
def _all_values_none(values: Iterable[object]) -> bool:
|
||||
return all(value is None for value in values)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS,
|
||||
UPSTREAM_USAGE_COST,
|
||||
|
|
@ -31,6 +32,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict
|
||||
)
|
||||
from litellm.litellm_core_utils.redact_messages import drop_token_ids
|
||||
from litellm.litellm_core_utils.token_counter import usage_tokenizer_model
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -434,6 +436,8 @@ class ChunkProcessor:
|
|||
chunks, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop"
|
||||
),
|
||||
message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""},
|
||||
provider_specific_fields=None,
|
||||
**_streamed_vllm_choice_fields(chain.from_iterable(chunk.get("choices") or () for chunk in chunks)),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -504,6 +508,7 @@ class ChunkProcessor:
|
|||
"index": 0,
|
||||
"message": {"role": role, "content": ""},
|
||||
"finish_reason": finish_reason,
|
||||
**_streamed_vllm_choice_fields(_first_choices(chunks)),
|
||||
}
|
||||
]
|
||||
),
|
||||
|
|
@ -1371,6 +1376,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),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1504,6 +1510,31 @@ def _field_of(value: object, name: str) -> object:
|
|||
return getattr(value, name, None)
|
||||
|
||||
|
||||
def _first_choices(chunks: Iterable[object]) -> tuple[object, ...]:
|
||||
return tuple(next(iter(object_items(_field_of(chunk, "choices"))), None) for chunk in chunks)
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -1776,6 +1807,19 @@ def stream_chunk_builder(
|
|||
end_time: object = None,
|
||||
logging_obj: Optional["Logging"] = None,
|
||||
count_prompt_tokens: Callable[[], int] | None = None,
|
||||
) -> ModelResponse | TextCompletionResponse | None:
|
||||
response: Final = _assembled_stream_response(chunks, messages, logging_obj, count_prompt_tokens)
|
||||
if response is not None and getattr(logging_obj, "drops_stream_token_ids", False) is True:
|
||||
for choice in response.choices:
|
||||
drop_token_ids(choice)
|
||||
return response
|
||||
|
||||
|
||||
def _assembled_stream_response(
|
||||
chunks: Sequence[object],
|
||||
messages: Sequence[AllMessageValues | Message] | None,
|
||||
logging_obj: Optional["Logging"],
|
||||
count_prompt_tokens: Callable[[], int] | None,
|
||||
) -> ModelResponse | TextCompletionResponse | None:
|
||||
try:
|
||||
if chunks is None:
|
||||
|
|
@ -1788,7 +1832,7 @@ def stream_chunk_builder(
|
|||
if not chunks:
|
||||
return None
|
||||
|
||||
processor: Final = ChunkProcessor(chunks, messages)
|
||||
processor: Final = ChunkProcessor(list(chunks), messages)
|
||||
chunks = processor.chunks
|
||||
|
||||
### BASE-CASE ###
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
UPSTREAM_USAGE_COST,
|
||||
|
|
@ -37,10 +38,11 @@ from litellm.litellm_core_utils.hidden_params import (
|
|||
set_hidden_param,
|
||||
stored_attribute,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_items
|
||||
from litellm.litellm_core_utils.model_response_utils import (
|
||||
is_model_response_stream_empty,
|
||||
)
|
||||
from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject
|
||||
from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject, drop_token_ids
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -137,6 +139,70 @@ def _prepend_upstream_choice_extras(
|
|||
choice.__pydantic_fields_set__.update(upstream_extras)
|
||||
|
||||
|
||||
_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_NO_CHOICE_FIELDS_BY_INDEX: Final[Mapping[int, 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 _response_choice_index(response_obj: Mapping[str, object]) -> int:
|
||||
choice_index: Final = response_obj.get("choice_index")
|
||||
if isinstance(choice_index, int):
|
||||
return choice_index
|
||||
return _upstream_choice_index(response_obj.get("original_chunk"))
|
||||
|
||||
|
||||
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: upstream_extras[field] for field in fields if upstream_extras.get(field) 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 _choice_fields_chunk(
|
||||
model_response: ModelResponseStream, upstream_index: int, fields: Mapping[str, object]
|
||||
) -> ModelResponseStream:
|
||||
model_response.choices[0].delta = Delta(content="") # 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], fields)
|
||||
return model_response
|
||||
|
||||
|
||||
def _without_token_ids_under_post_call_rules(chunk: ModelResponseStream | None) -> ModelResponseStream | None:
|
||||
if chunk is None or not litellm.post_call_rules:
|
||||
return chunk
|
||||
for choice in chunk.choices:
|
||||
drop_token_ids(choice)
|
||||
return chunk
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
|
@ -378,6 +444,8 @@ class CustomStreamWrapper:
|
|||
self._provider_response_model: str | None = None
|
||||
self._upstream_provider: str | None = None
|
||||
self.received_finish_reason: str | None = None
|
||||
self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX
|
||||
self.carried_chunk_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX
|
||||
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
|
||||
self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason
|
||||
self.special_tokens = [
|
||||
|
|
@ -823,6 +891,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
|
||||
|
|
@ -847,7 +916,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,11 +929,15 @@ class CustomStreamWrapper:
|
|||
is_finished = True
|
||||
finish_reason = choices[0].finish_reason
|
||||
usage = getattr(chunk, "usage", None)
|
||||
choice_index: Final = _upstream_choice_index(chunk)
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
"usage": usage,
|
||||
"chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS),
|
||||
"choice_index": choice_index,
|
||||
"sibling_finish_reason": finish_reason if choice_index != 0 else None,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -1183,6 +1256,13 @@ class CustomStreamWrapper:
|
|||
)
|
||||
|
||||
is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj)
|
||||
upstream_chunk_choice_fields: Final[Mapping[str, object]] = (
|
||||
_NO_UPSTREAM_CHOICE_FIELDS
|
||||
if litellm.post_call_rules
|
||||
else response_obj.get("chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS)
|
||||
)
|
||||
chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields)
|
||||
sibling_finish_reason: Final[str | None] = response_obj.get("sibling_finish_reason")
|
||||
|
||||
if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string
|
||||
self.raise_on_model_repetition()
|
||||
|
|
@ -1270,16 +1350,40 @@ class CustomStreamWrapper:
|
|||
if response_obj.get("provider_specific_fields") is not None:
|
||||
completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"]
|
||||
model_response.choices[0].delta = Delta(**completion_obj)
|
||||
_index: Final[int | None] = completion_obj.get("index")
|
||||
_index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index"))
|
||||
if _index is not None:
|
||||
model_response.choices[0].index = _index
|
||||
|
||||
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
|
||||
if sibling_finish_reason is not None:
|
||||
self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason)
|
||||
elif response_obj.get("is_finished"):
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
self._optional_combine_thinking_block_in_choices(model_response=model_response)
|
||||
|
||||
return model_response
|
||||
else:
|
||||
return
|
||||
elif sibling_finish_reason is not None:
|
||||
sibling_chunk: Final = _choice_fields_chunk(
|
||||
model_response, _response_choice_index(response_obj), chunk_choice_fields
|
||||
)
|
||||
self._finish_sibling_choice(sibling_chunk.choices[0], sibling_finish_reason)
|
||||
return sibling_chunk
|
||||
elif self.received_finish_reason is not None:
|
||||
chunk_upstream_index: Final = _response_choice_index(response_obj)
|
||||
if chunk_choice_fields and self._streams_past_the_pending_finish(
|
||||
model_response, response_obj, chunk_upstream_index
|
||||
):
|
||||
self.carried_chunk_choice_fields = MappingProxyType(
|
||||
{**self.carried_chunk_choice_fields, chunk_upstream_index: chunk_choice_fields}
|
||||
)
|
||||
else:
|
||||
another_choice_fields_chunk: Final = self._another_choice_fields_chunk(
|
||||
model_response, response_obj, chunk_choice_fields
|
||||
)
|
||||
if another_choice_fields_chunk is not None:
|
||||
return another_choice_fields_chunk
|
||||
if self.sent_last_chunk is True:
|
||||
# Bedrock returns the guardrail trace in the last chunk - we want to return this here
|
||||
if self.custom_llm_provider == "bedrock" and "trace" in model_response:
|
||||
|
|
@ -1317,12 +1421,20 @@ class CustomStreamWrapper:
|
|||
model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason(
|
||||
map_finish_reason(finish_reason=self.received_finish_reason)
|
||||
) # ensure consistent output to openai
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
if (
|
||||
response_obj.get("is_finished")
|
||||
and _response_choice_index(response_obj) == model_response.choices[0].index
|
||||
):
|
||||
_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 chunk_choice_fields:
|
||||
return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields)
|
||||
else:
|
||||
if hasattr(model_response, "usage"):
|
||||
self.chunks.append(model_response)
|
||||
|
|
@ -1545,6 +1657,8 @@ class CustomStreamWrapper:
|
|||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self._record_finish_choice_fields(chunk)
|
||||
if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj["usage"] is not None:
|
||||
_text_completion_usage: Final[Usage] = response_obj["usage"]
|
||||
|
|
@ -1608,6 +1722,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
|
||||
self.uses_openai_chat_stream = True
|
||||
if isinstance(chunk, ChatCompletionChunk):
|
||||
|
|
@ -1637,6 +1752,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._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)
|
||||
|
|
@ -1704,7 +1820,7 @@ class CustomStreamWrapper:
|
|||
completion_obj=completion_obj,
|
||||
)
|
||||
if isinstance(dispatch_result, _ProviderChunkEarlyReturn):
|
||||
return dispatch_result.value
|
||||
return _without_token_ids_under_post_call_rules(dispatch_result.value)
|
||||
response_obj = dispatch_result.response_obj
|
||||
|
||||
model_response.model = self._served_openrouter_model() or self.model
|
||||
|
|
@ -1794,7 +1910,7 @@ class CustomStreamWrapper:
|
|||
response_obj=response_obj,
|
||||
plain_delta_created=plain_delta_created,
|
||||
)
|
||||
return result
|
||||
return _without_token_ids_under_post_call_rules(result)
|
||||
|
||||
except StopIteration:
|
||||
raise StopIteration
|
||||
|
|
@ -1987,6 +2103,75 @@ class CustomStreamWrapper:
|
|||
return "tool_calls"
|
||||
return finish_reason
|
||||
|
||||
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_sibling_choice(self, choice: StreamingChoices, finish_reason: str) -> None:
|
||||
choice.finish_reason = map_finish_reason(finish_reason) # rebind-ok: chunks are built in place here
|
||||
self._apply_finish_choice_fields(choice)
|
||||
|
||||
def _streams_past_the_pending_finish(
|
||||
self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int
|
||||
) -> bool:
|
||||
finishing_index: Final = model_response.choices[0].index
|
||||
return (
|
||||
not self.sent_last_chunk
|
||||
and not response_obj.get("is_finished")
|
||||
and upstream_index != finishing_index
|
||||
and finishing_index in self.received_finish_choice_fields
|
||||
)
|
||||
|
||||
def _take_carried_chunk_choice_fields(
|
||||
self, response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
if not self.carried_chunk_choice_fields:
|
||||
return chunk_choice_fields
|
||||
upstream_index: Final = _response_choice_index(response_obj)
|
||||
carried: Final = self.carried_chunk_choice_fields.get(upstream_index)
|
||||
if carried is None:
|
||||
return chunk_choice_fields
|
||||
self.carried_chunk_choice_fields = MappingProxyType(
|
||||
{index: fields for index, fields in self.carried_chunk_choice_fields.items() if index != upstream_index}
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
field: [*object_items(carried.get(field)), *object_items(chunk_choice_fields.get(field))]
|
||||
for field in STREAM_CHUNK_CHOICE_FIELDS
|
||||
if field in carried or field in chunk_choice_fields
|
||||
}
|
||||
)
|
||||
|
||||
def _another_choice_fields_chunk(
|
||||
self,
|
||||
model_response: ModelResponseStream,
|
||||
response_obj: Mapping[str, object],
|
||||
chunk_choice_fields: Mapping[str, object],
|
||||
) -> ModelResponseStream | None:
|
||||
upstream_index: Final = _response_choice_index(response_obj)
|
||||
is_finished: Final = bool(response_obj.get("is_finished"))
|
||||
if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk:
|
||||
return None
|
||||
finish_fields: Final = (
|
||||
self.received_finish_choice_fields.get(upstream_index, _NO_UPSTREAM_CHOICE_FIELDS)
|
||||
if is_finished
|
||||
else _NO_UPSTREAM_CHOICE_FIELDS
|
||||
)
|
||||
if not chunk_choice_fields and all(value is None for value in finish_fields.values()):
|
||||
return None
|
||||
return _choice_fields_chunk(
|
||||
model_response, upstream_index, MappingProxyType({**finish_fields, **chunk_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}
|
||||
|
|
@ -2005,6 +2190,7 @@ class CustomStreamWrapper:
|
|||
model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason(
|
||||
_finish_reason if _finish_reason is not None else "stop"
|
||||
)
|
||||
self._apply_finish_choice_fields(model_response.choices[0])
|
||||
close_tag: Final = self._close_open_think_block()
|
||||
if close_tag is not None:
|
||||
model_response.choices[0].delta.content = close_tag
|
||||
|
|
|
|||
|
|
@ -493,10 +493,23 @@ class _CombinedChunkSplitter:
|
|||
for position, block in enumerate(blocks)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _without_sibling_finish(chunk: "ModelResponseStream") -> "ModelResponseStream":
|
||||
choices: Final = _optional_attr_sequence(chunk, "choices")
|
||||
if (
|
||||
len(choices) != 1
|
||||
or _optional_attr(choices[0], "index") in (None, 0)
|
||||
or _optional_attr(choices[0], "finish_reason") is None
|
||||
):
|
||||
return chunk
|
||||
unfinished: Final = copy.deepcopy(chunk)
|
||||
unfinished.choices[0].finish_reason = None
|
||||
return unfinished
|
||||
|
||||
def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]":
|
||||
split_chunks: Final = tuple(
|
||||
split_chunk
|
||||
for combined_chunk in self._split(chunk)
|
||||
for combined_chunk in self._split(self._without_sibling_finish(chunk))
|
||||
for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits
|
||||
)
|
||||
block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks)))
|
||||
|
|
|
|||
|
|
@ -90,6 +90,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.access_windows import access_windows_config_error
|
||||
from litellm.litellm_core_utils.add_retry_fallback_headers import (
|
||||
|
|
@ -10226,6 +10227,11 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -
|
|||
return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending
|
||||
|
||||
|
||||
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 _json_encodes(value: object) -> bool:
|
||||
try:
|
||||
to_jsonable_python(value)
|
||||
|
|
@ -10281,11 +10287,52 @@ def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes:
|
|||
|
||||
|
||||
def _serialize_streaming_chunk(chunk: BaseModel) -> bytes:
|
||||
choices: Final = _streaming_chunk_choices(chunk)
|
||||
null_choice_fields: Final = tuple(_null_vllm_choice_fields(choice) for choice in choices)
|
||||
if any(null_choice_fields):
|
||||
serialized_chunk: Final = _serialize_keeping_null_choice_fields(chunk, choices, null_choice_fields)
|
||||
if serialized_chunk is not None:
|
||||
return serialized_chunk
|
||||
if isinstance(chunk, ModelResponseStream):
|
||||
return _serialize_chat_chunk(chunk)
|
||||
return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode()
|
||||
|
||||
|
||||
def _serialize_keeping_null_choice_fields(
|
||||
chunk: BaseModel, choices: Sequence[BaseModel], null_choice_fields: Sequence[Mapping[str, None]]
|
||||
) -> bytes | None:
|
||||
try:
|
||||
return json.dumps(
|
||||
{
|
||||
**_dump_set_fields(chunk, exclude=frozenset({"choices"})),
|
||||
"choices": [
|
||||
{**_dump_set_fields(choice), **nulls}
|
||||
for choice, nulls in zip(choices, null_choice_fields, strict=True)
|
||||
],
|
||||
},
|
||||
separators=(",", ":"),
|
||||
ensure_ascii=False,
|
||||
).encode()
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _streaming_chunk_choices(chunk: BaseModel) -> Sequence[BaseModel]:
|
||||
match chunk:
|
||||
case ModelResponseStream() | TextCompletionResponse():
|
||||
return chunk.choices
|
||||
case _:
|
||||
return ()
|
||||
|
||||
|
||||
def _null_vllm_choice_fields(choice: BaseModel) -> Mapping[str, None]:
|
||||
return MappingProxyType({field: None for field, value in _vllm_choice_fields(choice).items() if value is None})
|
||||
|
||||
|
||||
def _dump_set_fields(model: BaseModel, exclude: AbstractSet[str] = frozenset()) -> Mapping[str, object]:
|
||||
return model.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=set(exclude))
|
||||
|
||||
|
||||
async def _apply_streaming_chunk_hooks(
|
||||
*,
|
||||
chunk: Any,
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
@ -170,6 +171,7 @@ from litellm.litellm_core_utils.litellm_logging import (
|
|||
deployment_pricing_model_info,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.litellm_core_utils.redact_messages import drop_token_ids
|
||||
from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
|
|
@ -179,6 +181,7 @@ from litellm.litellm_core_utils.served_output_texts import (
|
|||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy._types import (
|
||||
AlertType,
|
||||
|
|
@ -707,6 +710,56 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool:
|
|||
return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True)
|
||||
|
||||
|
||||
def _drop_stream_chunk_token_ids(chunk: object) -> None:
|
||||
for choice in stream_item_items(chunk, "choices"):
|
||||
drop_token_ids(choice)
|
||||
|
||||
|
||||
def _drop_built_stream_token_ids(request_data: Mapping[str, object]) -> None:
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if isinstance(logging_obj, Logging):
|
||||
logging_obj.drops_stream_token_ids = True
|
||||
|
||||
|
||||
def _drop_token_ids_ahead_of_chunk_check(
|
||||
callback: CustomLogger, chunk: object, request_data: Mapping[str, object]
|
||||
) -> None:
|
||||
if _overrides_hook(callback, "async_post_call_streaming_hook"):
|
||||
_drop_stream_chunk_token_ids(chunk)
|
||||
_drop_built_stream_token_ids(request_data)
|
||||
|
||||
|
||||
_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {})
|
||||
|
||||
|
||||
def _holds_stream_payload(value: object) -> bool:
|
||||
if isinstance(value, dict):
|
||||
nested: Final = _PLAIN_STREAM_FIELDS.validate_python(value).values()
|
||||
return any(item not in _EMPTY_STREAM_VALUES for item in nested)
|
||||
return value not in _EMPTY_STREAM_VALUES
|
||||
|
||||
|
||||
def _choice_carries_only_token_ids(choice: object) -> bool:
|
||||
if stream_item_field(choice, "token_ids") is None or not isinstance(choice, BaseModel):
|
||||
return False
|
||||
fields: Final = _PLAIN_STREAM_FIELDS.validate_python(choice.model_dump())
|
||||
return not any(_holds_stream_payload(value) for key, value in fields.items() if key not in ("index", "token_ids"))
|
||||
|
||||
|
||||
def _is_token_id_only_chunk(chunk: object) -> bool:
|
||||
choices: Final = stream_item_items(chunk, "choices")
|
||||
return bool(choices) and all(_choice_carries_only_token_ids(choice) for choice in choices)
|
||||
|
||||
|
||||
async def _without_token_id_only_chunks(response: AsyncIterable[object]) -> AsyncGenerator[object, None]:
|
||||
async for chunk in response:
|
||||
if _is_token_id_only_chunk(chunk):
|
||||
_drop_stream_chunk_token_ids(chunk)
|
||||
continue
|
||||
yield chunk
|
||||
|
||||
|
||||
def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool:
|
||||
return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines))
|
||||
|
||||
|
|
@ -4710,6 +4763,7 @@ class ProxyLogging:
|
|||
else:
|
||||
_callback = callback
|
||||
if _callback is not None and isinstance(_callback, CustomLogger):
|
||||
_drop_token_ids_ahead_of_chunk_check(_callback, response, data)
|
||||
if str_so_far is not None:
|
||||
complete_response = str_so_far + response_str
|
||||
else:
|
||||
|
|
@ -4731,7 +4785,7 @@ class ProxyLogging:
|
|||
self,
|
||||
response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_data: dict,
|
||||
request_data: dict[str, object],
|
||||
):
|
||||
"""
|
||||
Allow user to modify outgoing streaming data -> Given a whole response iterator.
|
||||
|
|
@ -4765,20 +4819,26 @@ class ProxyLogging:
|
|||
# Merge model-level guardrails before checking which guardrails to run
|
||||
request_data = check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router)
|
||||
|
||||
current_response = response
|
||||
stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
|
||||
|
||||
pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines)
|
||||
running_overrides: Final = tuple(
|
||||
(callback, kind)
|
||||
for callback, kind in caps.iterator_overrides
|
||||
if not isinstance(callback, CustomGuardrail)
|
||||
or (
|
||||
callback.guardrail_name not in pipeline_gated_names
|
||||
and callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is True
|
||||
)
|
||||
)
|
||||
drops_token_ids: Final = bool(post_call_pipelines) or any(
|
||||
isinstance(callback, CustomGuardrail) for callback, _ in running_overrides
|
||||
)
|
||||
if drops_token_ids:
|
||||
_drop_built_stream_token_ids(request_data)
|
||||
current_response = _without_token_id_only_chunks(response) if drops_token_ids else response
|
||||
guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect
|
||||
for resolved_callback, kind in caps.iterator_overrides:
|
||||
if isinstance(resolved_callback, CustomGuardrail):
|
||||
if resolved_callback.guardrail_name in pipeline_gated_names:
|
||||
continue
|
||||
if (
|
||||
resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
for resolved_callback, kind in running_overrides:
|
||||
effective_kind = (
|
||||
"apply_guardrail"
|
||||
if (
|
||||
|
|
@ -4836,6 +4896,8 @@ class ProxyLogging:
|
|||
try:
|
||||
async for chunk in current_response:
|
||||
served_chunks.append(chunk)
|
||||
if drops_token_ids:
|
||||
_drop_stream_chunk_token_ids(served_chunks[-1])
|
||||
yield chunk
|
||||
except (GeneratorExit, asyncio.CancelledError):
|
||||
await ProxyLogging._close_guarded_layers(guarded_layers)
|
||||
|
|
@ -10203,7 +10265,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,5 +1,6 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.constants import VLLM_CHOICE_FIELDS
|
||||
from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact
|
||||
from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse
|
||||
|
||||
|
|
@ -37,16 +38,19 @@ class TextCompletionStreamWrapper:
|
|||
response["object"] = "text_completion"
|
||||
response["created"] = chunk.get("created", None)
|
||||
response["model"] = chunk.get("model", None)
|
||||
text_choices: Final = TextChoices()
|
||||
if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices
|
||||
raise Exception
|
||||
|
||||
if chunk["choices"]:
|
||||
delta: Final = chunk["choices"][0]["delta"]
|
||||
chat_choice: Final = chunk["choices"][0]
|
||||
text_choices: Final = TextChoices(
|
||||
**{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice}
|
||||
)
|
||||
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"]
|
||||
response["choices"] = [text_choices]
|
||||
else:
|
||||
response["choices"] = []
|
||||
|
|
|
|||
|
|
@ -472,6 +472,30 @@ 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
|
||||
|
||||
|
||||
def test_cache_key_tells_apart_requests_that_ask_for_token_ids_through_extra_body() -> 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)
|
||||
other_extra_body_key: Final = cache.get_cache_key(**request, extra_body={"top_k": 5})
|
||||
|
||||
assert cache.get_cache_key(**request, extra_body={"return_token_ids": True}) != plain_key
|
||||
assert cache.get_cache_key(**request, extra_body={"top_k": 5, "return_token_ids": True}) != other_extra_body_key
|
||||
assert cache.get_cache_key(**request, extra_body={"return_token_ids": False}) == plain_key
|
||||
assert cache.get_cache_key(**request, extra_body={"continuation_token": "abc"}) != plain_key
|
||||
|
||||
|
||||
def test_cache_key_tells_apart_requests_that_differ_only_in_tool_call_id_shape() -> None:
|
||||
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
|
||||
request: Final = {
|
||||
|
|
|
|||
|
|
@ -3089,6 +3089,73 @@ def test_convert_to_model_response_object_function_output():
|
|||
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,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vllm_choice_fields",
|
||||
[{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": None, "token_ids": None}],
|
||||
ids=["stop-string-with-token-ids", "natural-stop"],
|
||||
)
|
||||
def test_vllm_choice_fields_cached_only_under_provider_specific_fields_come_back_on_the_choice(
|
||||
vllm_choice_fields: dict[str, object],
|
||||
) -> None:
|
||||
cached_before_upgrade: 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",
|
||||
"provider_specific_fields": vllm_choice_fields,
|
||||
}
|
||||
],
|
||||
}
|
||||
cache_hit: Final = convert_to_model_response_object(
|
||||
response_object=cached_before_upgrade, model_response_object=ModelResponse()
|
||||
)
|
||||
assert isinstance(cache_hit, ModelResponse)
|
||||
|
||||
dumped_choice: Final = json.loads(cache_hit.model_dump_json())["choices"][0]
|
||||
assert {field: dumped_choice.get(field, "missing") for field in vllm_choice_fields} == vllm_choice_fields
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"reported_cost",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -561,6 +561,83 @@ class TestPerformRedaction:
|
|||
{"text": "redacted-by-litellm", "finish_reason": "length", "index": 1},
|
||||
]
|
||||
|
||||
def test_redacts_token_ids_that_decode_to_the_completion(self):
|
||||
details = {
|
||||
"stream": True,
|
||||
"complete_streaming_response": litellm.ModelResponse(
|
||||
choices=[
|
||||
litellm.Choices(message=litellm.Message(content="secret", role="assistant"), token_ids=[20, 11])
|
||||
]
|
||||
),
|
||||
"standard_logging_object": {
|
||||
"response": {
|
||||
"choices": [
|
||||
{
|
||||
"token_ids": [20, 11],
|
||||
"message": {"content": "secret", "provider_specific_fields": {"token_ids": [20, 11]}},
|
||||
},
|
||||
{"text": " Paris.", "token_ids": [20, 11], "index": 1},
|
||||
]
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
perform_redaction(details, None)
|
||||
|
||||
logged_choices = details["standard_logging_object"]["response"]["choices"]
|
||||
assert [choice["token_ids"] for choice in logged_choices] == [None, None]
|
||||
assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None
|
||||
assert details["complete_streaming_response"].choices[0].get("token_ids") is None
|
||||
|
||||
def test_redacts_text_completion_stream_text_and_the_token_ids_that_decode_to_it(self):
|
||||
details = {
|
||||
"stream": True,
|
||||
"async_complete_streaming_response": litellm.TextCompletionResponse(
|
||||
choices=[
|
||||
litellm.utils.TextChoices(
|
||||
index=0, text=" Paris.", reasoning_content="thinking", token_ids=[20, 11], finish_reason="stop"
|
||||
)
|
||||
]
|
||||
),
|
||||
}
|
||||
|
||||
perform_redaction(details, None)
|
||||
|
||||
logged_choice = details["async_complete_streaming_response"].choices[0]
|
||||
assert (logged_choice.text, logged_choice.reasoning_content) == ("redacted-by-litellm", "redacted-by-litellm")
|
||||
assert logged_choice.get("token_ids") is None
|
||||
assert logged_choice.finish_reason == "stop"
|
||||
|
||||
def test_redacts_token_ids_a_cache_hit_keeps_on_the_choice_provider_specific_fields(self):
|
||||
result = litellm.ModelResponse(
|
||||
choices=[
|
||||
litellm.Choices(
|
||||
message=litellm.Message(content="secret", role="assistant"),
|
||||
provider_specific_fields={"token_ids": [20, 11], "stop_reason": "END"},
|
||||
)
|
||||
]
|
||||
)
|
||||
details = {
|
||||
"standard_logging_object": {
|
||||
"response": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {"content": "secret"},
|
||||
"provider_specific_fields": {"token_ids": [20, 11], "stop_reason": "END"},
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
redacted = perform_redaction(details, result)
|
||||
|
||||
logged_fields = details["standard_logging_object"]["response"]["choices"][0]["provider_specific_fields"]
|
||||
assert logged_fields == {"token_ids": None, "stop_reason": "END"}
|
||||
assert redacted.choices[0].provider_specific_fields["token_ids"] is None
|
||||
assert redacted.choices[0].provider_specific_fields["stop_reason"] == "END"
|
||||
assert result.choices[0].provider_specific_fields["token_ids"] == [20, 11]
|
||||
|
||||
def test_redacts_object_choices_inside_model_response_dict(self):
|
||||
result = {
|
||||
"choices": [
|
||||
|
|
|
|||
|
|
@ -2757,6 +2757,95 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it
|
|||
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_each_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 [
|
||||
{field: choice[field] for field in ("index", "stop_reason", "token_ids")}
|
||||
for choice in built.model_dump()["choices"]
|
||||
] == [
|
||||
{"index": 0, "stop_reason": ",5", "token_ids": [16, 11, 17, 11, 18, 19]},
|
||||
{"index": 1, "stop_reason": None, "token_ids": [22, 11]},
|
||||
]
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_the_vllm_fields_of_a_sole_choice_streamed_on_a_nonzero_index() -> None:
|
||||
built: Final = stream_chunk_builder(
|
||||
chunks=[
|
||||
_vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]),
|
||||
_vllm_chat_chunk(1, "8", "stop", stop_reason=",9", token_ids=[23]),
|
||||
]
|
||||
)
|
||||
|
||||
assert built is not None
|
||||
assert {field: built.model_dump()["choices"][0].get(field) for field in ("stop_reason", "token_ids")} == {
|
||||
"stop_reason": ",9",
|
||||
"token_ids": [22, 11, 23],
|
||||
}
|
||||
|
||||
|
||||
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])]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import threading
|
|||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import CodeType, FrameType, MappingProxyType, TracebackType
|
||||
from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast
|
||||
|
|
@ -27,6 +28,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,
|
||||
|
|
@ -34,6 +39,7 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
)
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator
|
||||
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.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator as GeminiModelResponseIterator,
|
||||
|
|
@ -6971,6 +6977,327 @@ 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)
|
||||
_CHOICE_INDEX: Final = TypeAdapter(int)
|
||||
_DELTA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
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(
|
||||
logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]]
|
||||
) -> CustomStreamWrapper:
|
||||
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 CustomStreamWrapper(
|
||||
completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj
|
||||
)
|
||||
|
||||
|
||||
def _hosted_vllm_stream_choices(
|
||||
logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]]
|
||||
) -> tuple[dict[str, object], ...]:
|
||||
return tuple(chunk.choices[0].model_dump() for chunk in _hosted_vllm_stream(logging_obj, upstream_choices))
|
||||
|
||||
|
||||
@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 [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason]
|
||||
assert [
|
||||
(_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"])
|
||||
for choice in emitted
|
||||
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)
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_sends_empty_content_on_token_id_only_chunks(logging_obj: Logging) -> 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": [11]},
|
||||
{"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [220]},
|
||||
{"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [16]},
|
||||
),
|
||||
)
|
||||
|
||||
assert [
|
||||
(_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids"))
|
||||
for choice in emitted
|
||||
if choice.get("token_ids")
|
||||
] == [("", [11]), ("", [220]), ("1", [16])]
|
||||
|
||||
|
||||
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 _rejects_a_two(text: str) -> bool:
|
||||
return "2" not in text
|
||||
|
||||
|
||||
def _choices_sent_before_rule_failure(stream: CustomStreamWrapper) -> Iterator[dict[str, object]]:
|
||||
try:
|
||||
yield from (chunk.choices[0].model_dump() for chunk in stream)
|
||||
except MidStreamFallbackError:
|
||||
return
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_sends_no_token_ids_ahead_of_a_failing_post_call_rule(
|
||||
logging_obj: Logging, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "post_call_rules", [_rejects_a_two])
|
||||
stream: Final = _hosted_vllm_stream(
|
||||
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"}, "finish_reason": None, "token_ids": [11, 17]},
|
||||
{"index": 0, "delta": {"content": ",2"}, "finish_reason": "stop", "stop_reason": None, "token_ids": [20]},
|
||||
),
|
||||
)
|
||||
|
||||
assert [
|
||||
(_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids"))
|
||||
for choice in _choices_sent_before_rule_failure(stream)
|
||||
] == [("", None), ("1", None)]
|
||||
|
||||
|
||||
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
|
||||
]
|
||||
|
||||
|
||||
_TWO_CHOICE_STOP_REASONS: Final[Mapping[int, str | int]] = MappingProxyType({0: ", 3", 1: 151645})
|
||||
_TWO_CHOICE_TOKEN_IDS: Final[Mapping[int, tuple[int, int]]] = MappingProxyType({0: (16, 17), 1: (26, 27)})
|
||||
|
||||
|
||||
def _two_choice_stream(first_finished: int, finish_contents: tuple[str, str]) -> tuple[Mapping[str, object], ...]:
|
||||
return (
|
||||
*({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)),
|
||||
*(
|
||||
{"index": index, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [ids[0]]}
|
||||
for index, ids in _TWO_CHOICE_TOKEN_IDS.items()
|
||||
),
|
||||
*(
|
||||
{
|
||||
"index": index,
|
||||
"delta": {"content": content},
|
||||
"finish_reason": "stop",
|
||||
"stop_reason": _TWO_CHOICE_STOP_REASONS[index],
|
||||
"token_ids": [_TWO_CHOICE_TOKEN_IDS[index][1]],
|
||||
}
|
||||
for index, content in zip((first_finished, 1 - first_finished), finish_contents)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _carries_another_choices_fields(choice: Mapping[str, object]) -> bool:
|
||||
own_index: Final = _CHOICE_INDEX.validate_python(choice["index"])
|
||||
token_ids: Final = _TOKEN_IDS.validate_python(choice.get("token_ids")) or ()
|
||||
return not set(token_ids) <= set(_TWO_CHOICE_TOKEN_IDS[own_index]) or choice.get("stop_reason") not in (
|
||||
None,
|
||||
_TWO_CHOICE_STOP_REASONS[own_index],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("first_finished", [0, 1])
|
||||
@pytest.mark.parametrize(
|
||||
"finish_contents",
|
||||
[("", ""), ("2", ""), ("2", "2")],
|
||||
ids=["both-finish-only", "content-then-finish-only", "both-content-then-finish"],
|
||||
)
|
||||
def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_own_choice(
|
||||
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))
|
||||
|
||||
assert [
|
||||
(choice["index"], choice.get("token_ids"), choice.get("stop_reason"))
|
||||
for choice in emitted
|
||||
if _carries_another_choices_fields(choice)
|
||||
] == []
|
||||
|
||||
|
||||
@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 _DELTA.validate_python(choice["delta"]).get("content") 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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("first_finished", [0, 1])
|
||||
@pytest.mark.parametrize(
|
||||
"finish_contents",
|
||||
[("", ""), ("2", ""), ("", "2"), ("2", "2")],
|
||||
ids=["both-finish-only", "content-then-finish-only", "finish-only-then-content", "both-content-then-finish"],
|
||||
)
|
||||
def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_reason_of_each_choice(
|
||||
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))
|
||||
|
||||
assert {
|
||||
index: _joined_token_ids([choice for choice in emitted if choice["index"] == index])
|
||||
for index in _TWO_CHOICE_TOKEN_IDS
|
||||
} == _TWO_CHOICE_TOKEN_IDS
|
||||
assert frozenset(
|
||||
(choice["index"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None
|
||||
) == frozenset(_TWO_CHOICE_STOP_REASONS.items())
|
||||
|
||||
|
||||
def test_hosted_vllm_stream_with_two_choices_finishes_the_first_choice_before_the_other_choice_streams_on(
|
||||
logging_obj: Logging,
|
||||
) -> None:
|
||||
emitted: Final = _hosted_vllm_stream_choices(
|
||||
logging_obj,
|
||||
(
|
||||
*({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)),
|
||||
{"index": 0, "delta": {"content": "1"}, "finish_reason": "stop", "stop_reason": ",2", "token_ids": [16]},
|
||||
{"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [26]},
|
||||
{"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]},
|
||||
{"index": 1, "delta": {"content": "1,"}, "finish_reason": None, "token_ids": [27]},
|
||||
{"index": 1, "delta": {"content": "3"}, "finish_reason": "stop", "stop_reason": 151645, "token_ids": [28]},
|
||||
),
|
||||
)
|
||||
first_choice_finish: Final = next(
|
||||
position for position, choice in enumerate(emitted) if choice["index"] == 0 and choice["finish_reason"]
|
||||
)
|
||||
|
||||
assert emitted[first_choice_finish]["stop_reason"] == ",2"
|
||||
assert _joined_token_ids([choice for choice in emitted[first_choice_finish:] if choice["index"] == 1]) == (
|
||||
26,
|
||||
11,
|
||||
27,
|
||||
28,
|
||||
)
|
||||
|
||||
|
||||
_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["finish_reason"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None
|
||||
] == [("stop", ", 5")]
|
||||
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, ...]:
|
||||
|
|
|
|||
|
|
@ -204,6 +204,20 @@ def test_splitter_splits_combined_chunk_into_content_then_finish():
|
|||
assert finish_chunk.choices[0].delta.content is None
|
||||
|
||||
|
||||
def test_a_finish_on_another_choice_index_does_not_end_the_message():
|
||||
chunks: Final = [
|
||||
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))]),
|
||||
ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content="Hi"))]),
|
||||
ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content=""), finish_reason="length")]),
|
||||
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]),
|
||||
]
|
||||
|
||||
events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="qwen3"))
|
||||
|
||||
stop_reasons: Final = [event["delta"]["stop_reason"] for event in events if event.get("type") == "message_delta"]
|
||||
assert stop_reasons == ["end_turn"]
|
||||
|
||||
|
||||
def test_is_combined_false_when_choices_empty():
|
||||
"""A metadata-only chunk with no choices is never treated as combined."""
|
||||
assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False
|
||||
|
|
|
|||
|
|
@ -3,14 +3,15 @@ import json
|
|||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
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():
|
||||
|
|
@ -401,6 +402,64 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema():
|
|||
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
|
||||
|
||||
|
||||
def _strict_weather_tool() -> dict[str, object]:
|
||||
return {
|
||||
"type": "function",
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import asyncio
|
|||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from collections.abc import AsyncIterator, Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -62,6 +62,8 @@ from litellm.types.utils import (
|
|||
Function,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
TextChoices,
|
||||
TextCompletionResponse,
|
||||
Usage,
|
||||
)
|
||||
from tests._support.clock_loops import AutojumpClockLoop
|
||||
|
|
@ -553,6 +555,102 @@ def test_serialize_streaming_chunk_simple_payload():
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("choice_extras", "serialized_extras"),
|
||||
[
|
||||
({"stop_reason": ",5", "token_ids": [20]}, {"stop_reason": ",5", "token_ids": [20]}),
|
||||
({"finish_reason": "stop", "stop_reason": None}, {"finish_reason": "stop", "stop_reason": None}),
|
||||
],
|
||||
ids=["vllm-fields", "natural-stop"],
|
||||
)
|
||||
def test_serialize_streaming_chunk_keeps_the_vllm_choice_fields(
|
||||
choice_extras: dict[str, object], serialized_extras: dict[str, object]
|
||||
):
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="chatcmpl-test",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content="1,2,3,4", role="assistant"), **choice_extras)],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [
|
||||
{"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras}
|
||||
]
|
||||
|
||||
|
||||
def _text_completion_chunk(choice_extras: Mapping[str, object]) -> TextCompletionResponse:
|
||||
return TextCompletionResponse(
|
||||
id="cmpl-1",
|
||||
choices=[TextChoices(finish_reason="length", index=0, text="", **choice_extras)],
|
||||
created=1,
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
)
|
||||
|
||||
|
||||
def _chat_chunk_with_fingerprint(choice_extras: Mapping[str, object]) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-1",
|
||||
choices=[StreamingChoices(finish_reason="length", index=0, delta=Delta(content=None), **choice_extras)],
|
||||
created=1,
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
object="chat.completion.chunk",
|
||||
system_fingerprint="fp_1",
|
||||
)
|
||||
|
||||
|
||||
_VLLM_SHAPED_CHUNKS: Final = pytest.mark.parametrize(
|
||||
"build_chunk", [_text_completion_chunk, _chat_chunk_with_fingerprint], ids=["text-completion", "chat-fingerprint"]
|
||||
)
|
||||
|
||||
|
||||
@_VLLM_SHAPED_CHUNKS
|
||||
@pytest.mark.parametrize(
|
||||
"vllm_choice_fields",
|
||||
[{"stop_reason": None}, {"stop_reason": None, "token_ids": [21]}, {"stop_reason": None, "token_ids": None}],
|
||||
ids=["null-stop-reason", "null-stop-reason-with-token-ids", "null-stop-reason-and-token-ids"],
|
||||
)
|
||||
def test_serialize_streaming_chunk_keeps_null_vllm_choice_fields(
|
||||
build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream],
|
||||
vllm_choice_fields: dict[str, object],
|
||||
) -> None:
|
||||
chunk: Final = build_chunk(vllm_choice_fields)
|
||||
plain_chunk: Final = json.loads(build_chunk({}).model_dump_json(exclude_none=True, exclude_unset=True))
|
||||
|
||||
assert json.loads(_serialize_streaming_chunk(chunk)) == {
|
||||
**plain_chunk,
|
||||
"choices": [{**plain_chunk["choices"][0], **vllm_choice_fields}],
|
||||
}
|
||||
|
||||
|
||||
@_VLLM_SHAPED_CHUNKS
|
||||
@pytest.mark.parametrize("choice_extras", [{}, {"stop_reason": ",5", "token_ids": [20]}], ids=["no-vllm", "set-vllm"])
|
||||
def test_serialize_streaming_chunk_without_null_vllm_fields_matches_model_dump_json(
|
||||
build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream],
|
||||
choice_extras: dict[str, object],
|
||||
) -> None:
|
||||
chunk: Final = build_chunk(choice_extras)
|
||||
|
||||
assert _serialize_streaming_chunk(chunk) == chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode()
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_keeps_streaming_a_null_stop_reason_chunk_with_a_value_json_cannot_encode() -> None:
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="chatcmpl-vllm",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content="a"), stop_reason=None, token_ids=[21])],
|
||||
created=0,
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
setattr(chunk, "trace_chunk", hashlib.sha256(b"trace").digest())
|
||||
|
||||
payload: Final = json.loads(_serialize_streaming_chunk(chunk))
|
||||
|
||||
assert payload["choices"][0]["delta"] == {"content": "a"}
|
||||
assert payload["choices"][0]["token_ids"] == [21]
|
||||
assert "trace_chunk" not in payload
|
||||
|
||||
|
||||
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``)."""
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -347,6 +347,241 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions
|
|||
assert delivered_text != ""
|
||||
|
||||
|
||||
async def _served_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, ...]:
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
||||
ContentFilterGuardrail,
|
||||
)
|
||||
from litellm.types.guardrails import BlockedWord, ContentFilterAction
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="output-filter",
|
||||
blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction("MASK"))],
|
||||
event_hook="post_call",
|
||||
default_on=guardrail_default_on,
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
|
||||
def chunk(content: str, token_ids: list[int], finish_reason: str | None = None) -> ModelResponseStream:
|
||||
streamed = ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)]
|
||||
)
|
||||
setattr(streamed.choices[0], "token_ids", token_ids)
|
||||
return streamed
|
||||
|
||||
async def fake_stream():
|
||||
yield chunk("the zebra runs", [1, 2, 3])
|
||||
yield chunk("", [4], finish_reason="stop")
|
||||
|
||||
served = [
|
||||
served_chunk
|
||||
async for served_chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
|
||||
response=fake_stream(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"),
|
||||
request_data={"model": "gpt-4o-mini", "metadata": {}},
|
||||
)
|
||||
]
|
||||
return tuple(getattr(served_chunk.choices[0], "token_ids", None) for served_chunk in served)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_guardrail_drops_token_ids_that_decode_to_the_unmasked_text(monkeypatch):
|
||||
served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=True)
|
||||
|
||||
assert served_token_ids
|
||||
assert all(token_ids is None for token_ids in served_token_ids)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_request(monkeypatch):
|
||||
served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=False)
|
||||
|
||||
assert served_token_ids == ([1, 2, 3], [4])
|
||||
|
||||
|
||||
class _ChecksEachChunkText(CustomGuardrail):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
|
||||
return None
|
||||
|
||||
|
||||
async def _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, bool]:
|
||||
import time
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[_ChecksEachChunkText(guardrail_name="chunk-check", event_hook="post_call", default_on=guardrail_default_on)],
|
||||
)
|
||||
logging_obj = Logging(
|
||||
model="m",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="per-chunk-call",
|
||||
function_id="per-chunk-call",
|
||||
)
|
||||
held_back_text_chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""))])
|
||||
setattr(held_back_text_chunk.choices[0], "token_ids", [220, 19, 11])
|
||||
|
||||
served = await ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_hook(
|
||||
data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj},
|
||||
response=held_back_text_chunk,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"),
|
||||
)
|
||||
return getattr(served.choices[0], "token_ids", None), logging_obj.drops_stream_token_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_chunk_stream_check_drops_token_ids_served_ahead_of_the_text_it_checks(monkeypatch):
|
||||
assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=True) == (None, True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_chunk_stream_keeps_token_ids_when_its_check_does_not_run_on_the_request(monkeypatch):
|
||||
assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=False) == ([220, 19, 11], False)
|
||||
|
||||
|
||||
class _ScansWholeStreamFirst(CustomGuardrail):
|
||||
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
||||
for chunk in [chunk async for chunk in response]:
|
||||
yield chunk
|
||||
|
||||
|
||||
async def _logged_stream_token_ids(monkeypatch, guardrail_default_on: bool) -> object:
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig
|
||||
|
||||
guardrail = _ScansWholeStreamFirst(
|
||||
guardrail_name="whole-stream-scan", event_hook="post_call", default_on=guardrail_default_on
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
logging_obj = Logging(
|
||||
model="m",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="token-ids-call",
|
||||
function_id="token-ids-call",
|
||||
)
|
||||
logged = asyncio.get_running_loop().create_future()
|
||||
|
||||
async def capture_logged_response(complete_streaming_response, cache_hit):
|
||||
logged.set_result(complete_streaming_response)
|
||||
|
||||
logging_obj._on_deferred_stream_complete = capture_logged_response
|
||||
|
||||
async def sse_lines():
|
||||
for choice in (
|
||||
{"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": ""}, "finish_reason": "stop", "token_ids": [20]},
|
||||
):
|
||||
yield "data: " + json.dumps(
|
||||
{"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [choice]}
|
||||
)
|
||||
yield "data: [DONE]"
|
||||
|
||||
stream = CustomStreamWrapper(
|
||||
completion_stream=HostedVLLMChatConfig().get_model_response_iterator(
|
||||
streaming_response=sse_lines(), sync_stream=False, json_mode=False
|
||||
),
|
||||
model="m",
|
||||
custom_llm_provider="hosted_vllm",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
async for _ in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
|
||||
response=stream,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"),
|
||||
request_data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj},
|
||||
):
|
||||
pass
|
||||
return (await asyncio.wait_for(logged, timeout=5)).choices[0].token_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_guardrail_drops_token_ids_from_the_response_it_logs_and_caches(monkeypatch):
|
||||
assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=True) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_logs_every_token_id_when_no_guardrail_runs_on_the_request(monkeypatch):
|
||||
assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=False) == [16, 11, 17, 20]
|
||||
|
||||
|
||||
def _completion_chunk(text: str, token_ids: list[int], finish_reason: str | None = None) -> object:
|
||||
from litellm.types.utils import TextChoices, TextCompletionResponse
|
||||
|
||||
return TextCompletionResponse(
|
||||
object="text_completion",
|
||||
choices=[TextChoices(index=0, text=text, finish_reason=finish_reason, token_ids=token_ids)],
|
||||
)
|
||||
|
||||
|
||||
def _guarded_completions_stream(monkeypatch, action: str, held_back: tuple[object, ...]):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
||||
ContentFilterGuardrail,
|
||||
)
|
||||
from litellm.types.guardrails import BlockedWord, ContentFilterAction
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[
|
||||
ContentFilterGuardrail(
|
||||
guardrail_name="output-filter",
|
||||
blocked_words=[BlockedWord(keyword="7", action=ContentFilterAction(action))],
|
||||
event_hook="post_call",
|
||||
default_on=True,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def upstream():
|
||||
for chunk in held_back:
|
||||
yield chunk
|
||||
yield _completion_chunk(", 6, 7", [21, 11, 22], finish_reason="stop")
|
||||
|
||||
return ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
|
||||
response=upstream(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/completions"),
|
||||
request_data={"model": "qwen3", "metadata": {}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_guardrail_blocks_completions_before_serving_token_id_only_chunks(monkeypatch):
|
||||
stream = _guarded_completions_stream(
|
||||
monkeypatch, "BLOCK", (_completion_chunk("", [11]), _completion_chunk("", [220]))
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await stream.__anext__()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_guardrail_serves_no_chunk_left_empty_by_the_token_id_drop(monkeypatch):
|
||||
held_back = (_completion_chunk("", [11]), _completion_chunk("", [220]))
|
||||
|
||||
served = [chunk async for chunk in _guarded_completions_stream(monkeypatch, "MASK", held_back)]
|
||||
|
||||
assert [(chunk.choices[0].text, chunk.choices[0].token_ids) for chunk in served] == [
|
||||
(", 6, [KEYWORD_REDACTED]", None)
|
||||
]
|
||||
assert [chunk.choices[0].token_ids for chunk in held_back] == [None, None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
|
|
|
|||
|
|
@ -13900,7 +13900,7 @@ def _stream_usage_generator_chunks():
|
|||
|
||||
def test_is_injected_stream_usage_artifact():
|
||||
from litellm.litellm_core_utils.stream_usage_chunks 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
|
||||
|
|
@ -13923,6 +13923,21 @@ 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
|
||||
|
||||
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,11 +1,15 @@
|
|||
from collections.abc import AsyncIterator, Iterator
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from itertools import chain
|
||||
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
|
||||
from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage
|
||||
|
||||
|
||||
def test_reasoning_content_preserved_in_text_completion_wrapper():
|
||||
|
|
@ -85,6 +89,146 @@ 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}],
|
||||
)
|
||||
|
||||
|
||||
def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]:
|
||||
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",
|
||||
)
|
||||
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"))
|
||||
|
||||
|
||||
@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",
|
||||
"index": 0,
|
||||
"stop_reason": ", 7",
|
||||
}
|
||||
|
||||
|
||||
def test_text_completion_stream_puts_stop_reason_on_the_final_text_chunk() -> None:
|
||||
responses: Final = _text_completion_stream(
|
||||
(_completion_chunk(", 5", None), _completion_chunk(", 6", "stop", stop_reason=", 7"))
|
||||
)
|
||||
|
||||
assert [(response.choices[0].get("text"), response.choices[0].get("stop_reason")) for response in responses] == [
|
||||
(", 5", None),
|
||||
(", 6", ", 7"),
|
||||
(None, ", 7"),
|
||||
]
|
||||
|
||||
|
||||
@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_stream_sends_empty_text_on_token_id_only_chunks() -> 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]),
|
||||
)
|
||||
)
|
||||
|
||||
choices: Final = [response.choices[0] for response in responses]
|
||||
|
||||
assert [(choice.text, choice.get("token_ids")) for choice in choices if choice.get("token_ids")] == [
|
||||
("", [11]),
|
||||
("", [220]),
|
||||
(", 5", [20]),
|
||||
]
|
||||
|
||||
|
||||
def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None:
|
||||
responses: Final = _text_completion_stream(
|
||||
(
|
||||
_completion_chunk(" 5", None, index=0, token_ids=[20]),
|
||||
_completion_chunk(" 7", None, index=1, token_ids=[22]),
|
||||
_completion_chunk(", 6", None, index=0, token_ids=[11, 21]),
|
||||
_completion_chunk("", None, index=1, token_ids=[11]),
|
||||
_completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]),
|
||||
_completion_chunk(", 8", "length", index=1, token_ids=[220, 23]),
|
||||
)
|
||||
)
|
||||
choices: Final = tuple(chain.from_iterable(response.choices for response in responses))
|
||||
|
||||
assert {
|
||||
index: (
|
||||
"".join(choice.text or "" for choice in choices if choice.index == index),
|
||||
list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)),
|
||||
)
|
||||
for index in (0, 1)
|
||||
} == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"upstream",
|
||||
[
|
||||
(
|
||||
_completion_chunk(" a", None, index=0),
|
||||
_completion_chunk(" b", None, index=1),
|
||||
_completion_chunk("", "length", index=2),
|
||||
_completion_chunk(" red", "stop", index=1, stop_reason=","),
|
||||
_completion_chunk(" s", "stop", index=0, stop_reason=","),
|
||||
),
|
||||
(
|
||||
_completion_chunk(" a", None, index=0),
|
||||
_completion_chunk(" b", None, index=1),
|
||||
_completion_chunk(" s", "stop", index=0, stop_reason=","),
|
||||
_completion_chunk("", "length", index=2),
|
||||
_completion_chunk(" red", "stop", index=1, stop_reason=","),
|
||||
),
|
||||
],
|
||||
ids=["other-choices-finish-first", "choice-0-finishes-first"],
|
||||
)
|
||||
def test_text_completion_stream_finishes_each_choice_on_its_own_index(upstream: Sequence[Completion]) -> None:
|
||||
finished: Final = [
|
||||
(choice.index, choice.finish_reason, choice.get("stop_reason"))
|
||||
for choice in chain.from_iterable(response.choices for response in _text_completion_stream(upstream))
|
||||
if choice.finish_reason is not None
|
||||
]
|
||||
|
||||
assert sorted(finished) == [(0, "stop", ","), (1, "stop", ","), (2, "length", None)]
|
||||
|
||||
|
||||
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