mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(streaming): keep finish_reason and provider usage when rebuilding streamed text completions (#44349)
* fix(streaming): keep finish_reason and provider usage when rebuilding streamed text completions stream_chunk_builder_text_completion read finish_reason from chunks[-1] and recounted usage from messages. With include_usage the last chunk is a usage-only trailer, so finish_reason came back None (or raised IndexError when the trailer had no choices), and prompt_tokens was 0 because a text-completion prompt is not in messages. Scan for the last finish_reason and prefer the provider-reported usage, as the chat path already does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(streaming): type the text completion chunk builder so its scans add no basedpyright debt Read chunk fields as typed attributes instead of untyped subscripts, which clears the reportUnknownArgumentType breach the lint gate flagged * fix(streaming): keep the provider's full usage, including token details, when rebuilding text completions Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(streaming): rebuild dict-form text completion chunks again stream_chunk_builder routes dict chunks holding TextChoices to the text completion builder, which now reads attributes. Convert them to TextCompletionResponse at the routing point. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(streaming): accept dict chunks in the text rebuild and derive its usage like the chat path Moves the dict to TextCompletionResponse conversion into stream_chunk_builder_text_completion so direct callers keep working, and builds usage through ChunkProcessor.calculate_usage over the chunks that report counts, so a trailer without total_tokens keeps its prompt and completion counts and the result is a fresh Usage * fix(streaming): read text completion usage through the shared usage extractor so usage dicts on typed chunks rebuild * fix(streaming): merge text completion logprobs across chunks and rebuild dict text chunks as text completions * fix(streaming): rebuild one text completion choice per index and keep null logprobs fields null * fix(streaming): keep dumped chat chunks with a choice-level text on the chat rebuild path Route a dict chunk to the text completion builder only when its choice has `text` and no `delta`. The Responses bridge snapshots chat chunks with model_dump(), and Together's chat choices carry `text` next to `delta`, so streamed /v1/responses calls to Together ended without response.completed. Group text completion choices by index in one sorted pass, so rebuilding an n>1 stream costs about the same as rebuilding an n=1 stream of the same size. * fix(streaming): keep text completion rebuilds working on mixed chunks and unreadable logprobs A chat chunk inside a text completion stream now converts to a text chunk that keeps its finish_reason and adds no text. A dict chunk whose logprobs fail validation is read again with logprobs dropped, and logprobs that are not a Logprobs object keep the last value sent. Each of these inputs raised or lost data at the previous commit, and each now rebuilds without raising Also adds a Purview test showing a multi-choice text stream is blocked like the non-streamed response, and Final annotations on the new test locals * fix(streaming): rebuild text completion choices without an index at index 0 Text completion choices whose index is None now rebuild under index 0, as on the merge base, where the sort by index raised at the previous commit. A chat chunk inside a text completion stream now keeps the usage it reports, and a choice whose logprobs mix a Logprobs object with another shape keeps the last value sent, as on the merge base * fix(streaming): rebuild text completion choices with a string index under that int A text completion choice whose index is a string now rebuilds under the int it names, and one whose index cannot be read as an int rebuilds at index 0, as on the merge base. A string index mixed with an int index raised TypeError at the previous commit, and a string-only index came back as a string * test(streaming): drop a docstring that repeats its test * fix(streaming): fall back to choice 0 for an infinite text completion index A float infinity index made int() raise OverflowError, which the index coercion did not catch, so the rebuild failed where the merge base returned choice 0. --------- Co-authored-by: Lai Wei <8022184+roywei@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Lai Wei
GitHub
parent
e34c779c67
commit
bda4e60f9b
6 changed files with 1093 additions and 53 deletions
|
|
@ -1,13 +1,13 @@
|
|||
import base64
|
||||
import builtins
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from itertools import chain, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Optional, TypeAlias, TypedDict, Union, cast
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeAlias, TypedDict, TypeVar, Union, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, Required
|
||||
|
||||
import litellm
|
||||
|
|
@ -17,6 +17,7 @@ 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.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantContentValue,
|
||||
ChatCompletionAudioDelta,
|
||||
)
|
||||
|
|
@ -36,11 +37,14 @@ from litellm.types.utils import (
|
|||
Function,
|
||||
FunctionCall,
|
||||
LlmProviders,
|
||||
Logprobs,
|
||||
Message,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
PromptTokensDetailsWrapper,
|
||||
ServerToolUse,
|
||||
StreamingChoices,
|
||||
TextChoices,
|
||||
TextCompletionResponse,
|
||||
Usage,
|
||||
)
|
||||
|
|
@ -857,7 +861,9 @@ class ChunkProcessor:
|
|||
return reasoning_tokens
|
||||
|
||||
@staticmethod
|
||||
def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None:
|
||||
def extract_usage_chunk(
|
||||
chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream | TextCompletionResponse",
|
||||
) -> Usage | None:
|
||||
usage_chunk: Usage | CompletionUsage | None = None
|
||||
if hasattr(chunk, "usage") and chunk.usage is not None:
|
||||
usage_chunk = chunk.usage
|
||||
|
|
@ -916,7 +922,7 @@ class ChunkProcessor:
|
|||
cost: float | None = None
|
||||
|
||||
for chunk in chunks:
|
||||
usage_chunk = self._extract_usage_chunk(chunk)
|
||||
usage_chunk = self.extract_usage_chunk(chunk)
|
||||
|
||||
if usage_chunk is not None:
|
||||
usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk)
|
||||
|
|
@ -1013,7 +1019,7 @@ class ChunkProcessor:
|
|||
values: Final = [
|
||||
value
|
||||
for chunk in chunks
|
||||
if (usage_chunk := self._extract_usage_chunk(chunk)) is not None
|
||||
if (usage_chunk := self.extract_usage_chunk(chunk)) is not None
|
||||
and isinstance(value := getattr(usage_chunk, field, None), str)
|
||||
]
|
||||
return values[-1] if values else None
|
||||
|
|
@ -1086,7 +1092,7 @@ class ChunkProcessor:
|
|||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
model: str,
|
||||
completion_output: str,
|
||||
messages: Sequence | None = None,
|
||||
messages: Sequence[object] | None = None,
|
||||
reasoning_tokens: int | None = None,
|
||||
count_prompt_tokens: Callable[[], int] | None = None,
|
||||
) -> Usage:
|
||||
|
|
@ -1199,7 +1205,7 @@ class ChunkProcessor:
|
|||
return returned_usage
|
||||
|
||||
|
||||
def concatenate_base64_list(base64_strings: list[str]) -> str:
|
||||
def concatenate_base64_list(base64_strings: Sequence[str]) -> str:
|
||||
"""
|
||||
Concatenates a list of base64-encoded strings.
|
||||
|
||||
|
|
@ -1216,35 +1222,115 @@ def concatenate_base64_list(base64_strings: list[str]) -> str:
|
|||
return base64.b64encode(combined_bytes).decode("utf-8")
|
||||
|
||||
|
||||
def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse:
|
||||
id: Final = chunks[0]["id"]
|
||||
object: Final = chunks[0]["object"]
|
||||
created: Final = chunks[0]["created"]
|
||||
model: Final = chunks[0]["model"]
|
||||
system_fingerprint: Final = chunks[0].get("system_fingerprint", None)
|
||||
finish_reason: Final = chunks[-1]["choices"][0]["finish_reason"]
|
||||
logprobs: Final = chunks[-1]["choices"][0]["logprobs"]
|
||||
_CHOICE_DICTS: Final = TypeAdapter(list[dict[str, object]])
|
||||
|
||||
content_list: Final = []
|
||||
for chunk in chunks:
|
||||
choices = chunk["choices"]
|
||||
for choice in choices:
|
||||
if choice is not None and hasattr(choice, "text") and choice.get("text") is not None:
|
||||
_choice = choice.get("text")
|
||||
content_list.append(_choice)
|
||||
|
||||
# Combine the "content" strings into a single string || combine the 'function' strings into a single string
|
||||
combined_content: Final = "".join(content_list)
|
||||
|
||||
def _text_completion_chunk_from_mapping(chunk: Mapping[str, object]) -> TextCompletionResponse:
|
||||
try:
|
||||
prompt_tokens = token_counter(model=model, messages=messages)
|
||||
except Exception: # don't allow this failing to block a complete streaming response from being returned
|
||||
print_verbose("token_counter failed, assuming prompt tokens is 0")
|
||||
prompt_tokens = 0
|
||||
completion_tokens: Final = token_counter(
|
||||
model=model,
|
||||
text=combined_content,
|
||||
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
|
||||
return TextCompletionResponse.model_validate(chunk)
|
||||
except ValidationError:
|
||||
choices: Final = _CHOICE_DICTS.validate_python(chunk.get("choices", []))
|
||||
return TextCompletionResponse.model_validate(
|
||||
{**chunk, "choices": [{**choice, "logprobs": None} for choice in choices]}
|
||||
)
|
||||
|
||||
|
||||
def _as_text_completion_chunk(
|
||||
chunk: TextCompletionResponse | ModelResponseStream | Mapping[str, object],
|
||||
) -> TextCompletionResponse:
|
||||
if isinstance(chunk, TextCompletionResponse):
|
||||
return chunk
|
||||
if isinstance(chunk, ModelResponseStream):
|
||||
return TextCompletionResponse(
|
||||
id=chunk.id,
|
||||
created=chunk.created,
|
||||
model=chunk.model,
|
||||
choices=[TextChoices(index=choice.index, finish_reason=choice.finish_reason) for choice in chunk.choices],
|
||||
usage=ChunkProcessor.extract_usage_chunk(chunk),
|
||||
)
|
||||
return _text_completion_chunk_from_mapping(chunk)
|
||||
|
||||
|
||||
_LogprobsValue: Final = TypeVar("_LogprobsValue")
|
||||
|
||||
|
||||
def _joined_or_none(fields: Iterable[Sequence[_LogprobsValue] | None]) -> list[_LogprobsValue] | None:
|
||||
sent: Final = tuple(field for field in fields if field is not None)
|
||||
return list(chain.from_iterable(sent)) if sent else None
|
||||
|
||||
|
||||
class _CarriesLogprobs(Protocol):
|
||||
@property
|
||||
def logprobs(self) -> object: ...
|
||||
|
||||
|
||||
def _sent_logprobs(choice: _CarriesLogprobs) -> object:
|
||||
return choice.logprobs
|
||||
|
||||
|
||||
def _merged_logprobs(choices: Sequence[TextChoices]) -> object:
|
||||
sent: Final = tuple(logprobs for logprobs in map(_sent_logprobs, choices) if logprobs is not None)
|
||||
parts: Final = tuple(logprobs for logprobs in sent if isinstance(logprobs, Logprobs))
|
||||
if len(parts) < len(sent):
|
||||
return sent[-1]
|
||||
if not parts:
|
||||
return None
|
||||
return Logprobs(
|
||||
text_offset=_joined_or_none(part.text_offset for part in parts),
|
||||
token_logprobs=_joined_or_none(part.token_logprobs for part in parts),
|
||||
tokens=_joined_or_none(part.tokens for part in parts),
|
||||
top_logprobs=_joined_or_none(part.top_logprobs for part in parts),
|
||||
)
|
||||
|
||||
|
||||
def _text_of(choice: TextChoices) -> str:
|
||||
return choice.text if isinstance(choice.text, str) else ""
|
||||
|
||||
|
||||
def _index_of(choice: TextChoices) -> int:
|
||||
try:
|
||||
return int(choice.index or 0)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return 0
|
||||
|
||||
|
||||
def _rebuilt_text_choice(index: int, choices: Sequence[TextChoices]) -> TextChoices:
|
||||
return TextChoices(
|
||||
text="".join(map(_text_of, choices)),
|
||||
index=index,
|
||||
logprobs=_merged_logprobs(choices),
|
||||
finish_reason=next((choice.finish_reason for choice in reversed(choices) if choice.finish_reason), None),
|
||||
)
|
||||
|
||||
|
||||
def _reported_usage(chunk: TextCompletionResponse) -> Usage | None:
|
||||
usage: Final = ChunkProcessor.extract_usage_chunk(chunk)
|
||||
return usage if usage is not None and (usage.prompt_tokens or usage.completion_tokens) else None
|
||||
|
||||
|
||||
def stream_chunk_builder_text_completion(
|
||||
chunks: Sequence[TextCompletionResponse | ModelResponseStream | Mapping[str, object]],
|
||||
messages: Sequence[AllMessageValues | Message] | None = None,
|
||||
) -> TextCompletionResponse:
|
||||
text_chunks: Final = tuple(_as_text_completion_chunk(chunk) for chunk in chunks)
|
||||
id: Final = text_chunks[0].id
|
||||
object: Final = text_chunks[0].object
|
||||
created: Final = text_chunks[0].created
|
||||
model: Final = text_chunks[0].model
|
||||
system_fingerprint: Final = text_chunks[0].get("system_fingerprint", None)
|
||||
all_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in text_chunks))
|
||||
choices: Final = [
|
||||
_rebuilt_text_choice(index, tuple(group))
|
||||
for index, group in groupby(sorted(all_choices, key=_index_of), key=_index_of)
|
||||
] or [_rebuilt_text_choice(0, ())]
|
||||
combined_content: Final = "".join(map(_text_of, choices))
|
||||
|
||||
reported_usages: Final = tuple(reported for reported in map(_reported_usage, text_chunks) if reported is not None)
|
||||
usage: Final = ChunkProcessor(list(text_chunks), messages).calculate_usage(
|
||||
chunks=tuple(_UsageBearingChunk(usage=reported) for reported in reported_usages),
|
||||
model=model or "",
|
||||
completion_output=combined_content,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
response: Final = {
|
||||
|
|
@ -1253,19 +1339,8 @@ def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None
|
|||
"created": created,
|
||||
"model": model,
|
||||
"system_fingerprint": system_fingerprint,
|
||||
"choices": [
|
||||
{
|
||||
"text": combined_content,
|
||||
"index": 0,
|
||||
"logprobs": logprobs,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
"choices": choices,
|
||||
"usage": usage,
|
||||
}
|
||||
return TextCompletionResponse(**response)
|
||||
|
||||
|
|
@ -1416,9 +1491,14 @@ def stream_chunk_builder(
|
|||
return None
|
||||
## Route to the text completion logic
|
||||
first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), None)
|
||||
if first_chunk_with_choices is not None and isinstance(
|
||||
first_chunk_with_choices["choices"][0], litellm.utils.TextChoices
|
||||
): # route to the text completion logic
|
||||
if first_chunk_with_choices is not None and (
|
||||
isinstance(first_chunk_with_choices["choices"][0], TextChoices)
|
||||
or (
|
||||
isinstance(first_chunk_with_choices["choices"][0], Mapping)
|
||||
and "text" in first_chunk_with_choices["choices"][0]
|
||||
and "delta" not in first_chunk_with_choices["choices"][0]
|
||||
)
|
||||
):
|
||||
return stream_chunk_builder_text_completion(chunks=chunks, messages=messages)
|
||||
|
||||
model: Final = chunks[0]["model"]
|
||||
|
|
|
|||
|
|
@ -9166,7 +9166,7 @@ from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( # noqa:
|
|||
stream_chunk_builder as stream_chunk_builder, # noqa: PLC0414 # pyright: ignore[reportUnknownVariableType] # chunks is an untyped list
|
||||
)
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( # noqa: E402 # re-export after the move
|
||||
stream_chunk_builder_text_completion as stream_chunk_builder_text_completion, # noqa: PLC0414 # pyright: ignore[reportUnknownVariableType] # chunks is an untyped list
|
||||
stream_chunk_builder_text_completion as stream_chunk_builder_text_completion, # noqa: PLC0414 # re-export after the move
|
||||
)
|
||||
|
||||
########## Token Counting API ##########
|
||||
|
|
|
|||
280
tests/integration/sdk/test_text_completion_stream_rebuild_sdk.py
Normal file
280
tests/integration/sdk/test_text_completion_stream_rebuild_sdk.py
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.timing import assert_within_factor_of_control, timed
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import TextCompletionResponse
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_PROVIDER_KEY: Final = "synthetic-provider-key"
|
||||
_PIECES: Final = ("Hello", " streamed", " world")
|
||||
_PROMPT_TOKENS: Final = 11
|
||||
_COMPLETION_TOKENS: Final = 3
|
||||
_CACHED_TOKENS: Final = 4
|
||||
_USAGE: Final[dict[str, JsonValue]] = {
|
||||
"prompt_tokens": _PROMPT_TOKENS,
|
||||
"completion_tokens": _COMPLETION_TOKENS,
|
||||
"total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
|
||||
"prompt_tokens_details": {"cached_tokens": _CACHED_TOKENS},
|
||||
}
|
||||
_CHOICE_PIECES: Final = (("Hi", " there"), ("Hey", " you"))
|
||||
_GRID_CHOICES: Final = 128
|
||||
_GRID_CHUNKS: Final = 8192
|
||||
_GRID_ROUNDS: Final = 2
|
||||
_REBUILD_FACTOR: Final = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Member:
|
||||
model: str
|
||||
target: str
|
||||
|
||||
|
||||
_MEMBERS: Final = {
|
||||
"text_completion_openai": Member("text-completion-openai/gpt-3.5-turbo-instruct", "/v1/completions"),
|
||||
"hosted_vllm": Member("hosted_vllm/Qwen/Qwen3-0.6B", "/v1/completions"),
|
||||
"openai_chat_backed": Member("openai/gpt-5.4-nano", "/v1/chat/completions"),
|
||||
}
|
||||
|
||||
|
||||
def _frame(body: Mapping[str, JsonValue]) -> bytes:
|
||||
return b"data: " + json.dumps(body).encode() + b"\n\n"
|
||||
|
||||
|
||||
def _text_choice(index: int, text: str, finish_reason: str | None) -> dict[str, JsonValue]:
|
||||
return {"text": text, "index": index, "logprobs": None, "finish_reason": finish_reason}
|
||||
|
||||
|
||||
def _chat_choice(index: int, text: str, finish_reason: str | None) -> dict[str, JsonValue]:
|
||||
return {"index": index, "delta": {"role": "assistant", "content": text}, "finish_reason": finish_reason}
|
||||
|
||||
|
||||
def _stream(
|
||||
request: Request, identity: str, pieces: tuple[tuple[int, str, str | None], ...], include_usage: bool
|
||||
) -> Reply:
|
||||
chat: Final = request.target.endswith("/chat/completions")
|
||||
head: Final[dict[str, JsonValue]] = {
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk" if chat else "text_completion",
|
||||
"created": 1,
|
||||
"model": "scripted",
|
||||
}
|
||||
choice: Final = _chat_choice if chat else _text_choice
|
||||
usage: Final = (_frame({**head, "choices": [], "usage": _USAGE}),) if include_usage else ()
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
*(_frame({**head, "choices": [choice(index, text, finish)]}) for index, text, finish in pieces),
|
||||
*usage,
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _asks_for_usage(body: Mapping[str, JsonValue]) -> bool:
|
||||
options: Final = body.get("stream_options")
|
||||
return isinstance(options, dict) and options.get("include_usage") is True
|
||||
|
||||
|
||||
def _upstream(identity: str, pieces: tuple[tuple[int, str, str | None], ...]) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
assert body.get("stream") is True, body
|
||||
return _stream(request, identity, pieces, _asks_for_usage(body))
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _one_choice_pieces() -> tuple[tuple[int, str, str | None], ...]:
|
||||
return (
|
||||
*((0, piece, None) for piece in _PIECES),
|
||||
(0, "", "length"),
|
||||
)
|
||||
|
||||
|
||||
def _two_choice_pieces() -> tuple[tuple[int, str, str | None], ...]:
|
||||
(first, second) = _CHOICE_PIECES
|
||||
return ((0, first[0], None), (1, second[0], None), (0, first[1], "stop"), (1, second[1], "length"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def identity() -> str:
|
||||
return "cmpl-rebuild-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def one_choice_wire(identity: str) -> Iterator[Wire]:
|
||||
with wire_server(_upstream(identity, _one_choice_pieces())) as served:
|
||||
yield served
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def two_choice_wire(identity: str) -> Iterator[Wire]:
|
||||
with wire_server(_upstream(identity, _two_choice_pieces())) as served:
|
||||
yield served
|
||||
|
||||
|
||||
def _stream_options(include_usage: bool) -> dict[str, object]:
|
||||
return {"stream_options": {"include_usage": True}} if include_usage else {}
|
||||
|
||||
|
||||
async def _async_chunks(
|
||||
member: Member, wire: Wire, include_usage: bool, extra: Mapping[str, object]
|
||||
) -> tuple[TextCompletionResponse, ...]:
|
||||
stream: Final = await litellm.atext_completion(
|
||||
model=member.model,
|
||||
prompt="rebuild probe",
|
||||
api_base=wire.url + "/v1",
|
||||
api_key=_PROVIDER_KEY,
|
||||
max_tokens=16,
|
||||
stream=True,
|
||||
num_retries=0,
|
||||
**_stream_options(include_usage),
|
||||
**extra,
|
||||
)
|
||||
return tuple([chunk async for chunk in stream])
|
||||
|
||||
|
||||
def _sync_chunks(member: Member, wire: Wire, include_usage: bool) -> tuple[TextCompletionResponse, ...]:
|
||||
stream: Final = litellm.text_completion(
|
||||
model=member.model,
|
||||
prompt="rebuild probe",
|
||||
api_base=wire.url + "/v1",
|
||||
api_key=_PROVIDER_KEY,
|
||||
max_tokens=16,
|
||||
stream=True,
|
||||
num_retries=0,
|
||||
**_stream_options(include_usage),
|
||||
)
|
||||
return tuple(stream)
|
||||
|
||||
|
||||
def _rebuilt(chunks: tuple[TextCompletionResponse, ...]) -> TextCompletionResponse:
|
||||
rebuilt: Final = litellm.stream_chunk_builder(chunks=list(chunks))
|
||||
assert isinstance(rebuilt, TextCompletionResponse), rebuilt
|
||||
return rebuilt
|
||||
|
||||
|
||||
def _posted_targets(wire: Wire) -> tuple[str, ...]:
|
||||
return tuple(request.target for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("member", "mode"),
|
||||
[
|
||||
("text_completion_openai", "async"),
|
||||
("text_completion_openai", "sync"),
|
||||
("hosted_vllm", "async"),
|
||||
("openai_chat_backed", "async"),
|
||||
],
|
||||
)
|
||||
async def test_rebuilt_text_completion_stream_keeps_the_reported_finish_reason_and_usage(
|
||||
one_choice_wire: Wire, identity: str, member: str, mode: str
|
||||
) -> None:
|
||||
chosen: Final = _MEMBERS[member]
|
||||
chunks: Final = (
|
||||
await _async_chunks(chosen, one_choice_wire, True, {})
|
||||
if mode == "async"
|
||||
else _sync_chunks(chosen, one_choice_wire, True)
|
||||
)
|
||||
|
||||
rebuilt: Final = _rebuilt(chunks)
|
||||
|
||||
assert [(choice.index, choice.text) for choice in rebuilt.choices] == [(0, "".join(_PIECES))], rebuilt
|
||||
assert rebuilt.choices[0].finish_reason == "length", rebuilt
|
||||
assert (rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens, rebuilt.usage.total_tokens) == (
|
||||
_PROMPT_TOKENS,
|
||||
_COMPLETION_TOKENS,
|
||||
_PROMPT_TOKENS + _COMPLETION_TOKENS,
|
||||
), rebuilt.usage
|
||||
assert rebuilt.usage.prompt_tokens_details is not None, rebuilt.usage
|
||||
assert rebuilt.usage.prompt_tokens_details.cached_tokens == _CACHED_TOKENS, rebuilt.usage
|
||||
assert rebuilt.id == identity, rebuilt
|
||||
assert _posted_targets(one_choice_wire) == (chosen.target,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("member", ["text_completion_openai", "openai_chat_backed"])
|
||||
async def test_rebuilt_text_completion_stream_without_reported_usage_counts_the_streamed_text(
|
||||
one_choice_wire: Wire, member: str
|
||||
) -> None:
|
||||
chosen: Final = _MEMBERS[member]
|
||||
chunks: Final = await _async_chunks(chosen, one_choice_wire, False, {})
|
||||
|
||||
rebuilt: Final = _rebuilt(chunks)
|
||||
|
||||
text: Final = "".join(_PIECES)
|
||||
assert rebuilt.choices[0].text == text, rebuilt
|
||||
assert rebuilt.choices[0].finish_reason == "length", rebuilt
|
||||
expected_completion: Final = litellm.token_counter(model=chosen.model, text=text, count_response_tokens=True)
|
||||
assert expected_completion > 0
|
||||
assert rebuilt.usage.completion_tokens == expected_completion, rebuilt.usage
|
||||
assert rebuilt.usage.total_tokens == rebuilt.usage.prompt_tokens + rebuilt.usage.completion_tokens, rebuilt.usage
|
||||
assert _posted_targets(one_choice_wire) == (chosen.target,)
|
||||
|
||||
|
||||
async def test_rebuilt_chat_backed_text_completion_stream_keeps_each_choice_text_apart(two_choice_wire: Wire) -> None:
|
||||
chosen: Final = _MEMBERS["openai_chat_backed"]
|
||||
chunks: Final = await _async_chunks(chosen, two_choice_wire, True, {"n": 2})
|
||||
|
||||
rebuilt: Final = _rebuilt(chunks)
|
||||
|
||||
assert [(choice.index, choice.text) for choice in rebuilt.choices] == [
|
||||
(index, "".join(pieces)) for index, pieces in enumerate(_CHOICE_PIECES)
|
||||
], rebuilt
|
||||
assert (rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) == (_PROMPT_TOKENS, _COMPLETION_TOKENS)
|
||||
assert _posted_targets(two_choice_wire) == (chosen.target,)
|
||||
|
||||
|
||||
def _grid_row(step: int, last: int, choices: int) -> Iterator[tuple[int, str, str | None]]:
|
||||
return ((index, "x", "length" if step == last else None) for index in range(choices))
|
||||
|
||||
|
||||
def _grid_pieces(choices: int) -> tuple[tuple[int, str, str | None], ...]:
|
||||
steps: Final = _GRID_CHUNKS // choices
|
||||
return tuple(chain.from_iterable(_grid_row(step, steps - 1, choices) for step in range(steps)))
|
||||
|
||||
|
||||
def _grid_upstream(request: Request) -> Reply:
|
||||
requested: Final = _JSON_OBJECT.validate_json(request.body).get("n")
|
||||
choices: Final = requested if isinstance(requested, int) else 1
|
||||
return _stream(request, "cmpl-grid-" + uuid.uuid4().hex, _grid_pieces(choices), False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def grid_wire() -> Iterator[Wire]:
|
||||
with wire_server(_grid_upstream) as served:
|
||||
yield served
|
||||
|
||||
|
||||
async def _rebuild_seconds(wire: Wire, choices: int) -> float:
|
||||
chunks: Final = await _async_chunks(_MEMBERS["openai_chat_backed"], wire, False, {"n": choices})
|
||||
run: Final = timed(lambda: _rebuilt(chunks))
|
||||
assert sum(len(choice.text) for choice in run.value.choices) == _GRID_CHUNKS, run.value.usage
|
||||
return run.seconds
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
async def test_rebuilding_a_many_choice_text_completion_stream_costs_what_one_choice_of_the_same_size_does(
|
||||
grid_wire: Wire, record_property: Callable[[str, object], None]
|
||||
) -> None:
|
||||
pairs: Final = [
|
||||
(await _rebuild_seconds(grid_wire, 1), await _rebuild_seconds(grid_wire, _GRID_CHOICES))
|
||||
for _ in range(_GRID_ROUNDS)
|
||||
]
|
||||
record_property("one_choice_rebuild_seconds", [control for control, _ in pairs])
|
||||
record_property("many_choice_rebuild_seconds", [observed for _, observed in pairs])
|
||||
|
||||
assert_within_factor_of_control(
|
||||
min(observed for _, observed in pairs),
|
||||
control=min(control for control, _ in pairs),
|
||||
factor=_REBUILD_FACTOR,
|
||||
)
|
||||
assert len(_posted_targets(grid_wire)) == 2 * _GRID_ROUNDS
|
||||
|
|
@ -0,0 +1,286 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
list_value,
|
||||
object_value,
|
||||
)
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
|
||||
pytestmark: Final = pytest.mark.timeout(_CELL_SECONDS)
|
||||
|
||||
_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_PROVIDER_KEY: Final = "synthetic-provider-key"
|
||||
_GUARDRAIL: Final = "audit-text-rebuild-tool-permission"
|
||||
_TEXT_MODEL: Final = f"audit-text-instruct-{uuid.uuid4().hex}"
|
||||
_PIECES: Final = ("Hello", " streamed", " world")
|
||||
_PROMPT_TOKENS: Final = 11
|
||||
_COMPLETION_TOKENS: Final = 3
|
||||
_USAGE: Final[dict[str, JsonValue]] = {
|
||||
"prompt_tokens": _PROMPT_TOKENS,
|
||||
"completion_tokens": _COMPLETION_TOKENS,
|
||||
"total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
|
||||
}
|
||||
_BRIDGE_TEXT: Final = "bridge reply"
|
||||
_TOGETHER_MODEL: Final = "together_ai/Qwen/Qwen3.5-9B"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
wire: Wire
|
||||
proxy: OwnedProxy
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Stream:
|
||||
status: int
|
||||
body: bytes
|
||||
|
||||
@property
|
||||
def frames(self) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
JSON_OBJECT.validate_json(line[6:])
|
||||
for line in self.body.decode().splitlines()
|
||||
if line.startswith("data: {")
|
||||
)
|
||||
|
||||
@property
|
||||
def choices(self) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(chain.from_iterable(_choices_of(frame) for frame in self.frames))
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
return "".join(str(choice.get("text") or "") for choice in self.choices)
|
||||
|
||||
@property
|
||||
def finish_reasons(self) -> tuple[JsonValue, ...]:
|
||||
return tuple(choice["finish_reason"] for choice in self.choices if choice.get("finish_reason") is not None)
|
||||
|
||||
@property
|
||||
def usages(self) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(object_value(frame["usage"]) for frame in self.frames if frame.get("usage") is not None)
|
||||
|
||||
@property
|
||||
def ids(self) -> frozenset[JsonValue]:
|
||||
return frozenset(frame["id"] for frame in self.frames)
|
||||
|
||||
|
||||
def _choices_of(frame: Mapping[str, JsonValue]) -> Iterator[Mapping[str, JsonValue]]:
|
||||
return (object_value(choice) for choice in list_value(frame.get("choices") or []))
|
||||
|
||||
|
||||
def _frame(body: Mapping[str, JsonValue]) -> bytes:
|
||||
return b"data: " + json.dumps(body).encode() + b"\n\n"
|
||||
|
||||
|
||||
def _text_frame(identity: str, choices: list[JsonValue], **extra: JsonValue) -> bytes:
|
||||
return _frame(
|
||||
{"id": identity, "object": "text_completion", "created": 1, "model": "scripted", "choices": choices, **extra}
|
||||
)
|
||||
|
||||
|
||||
def _text_choice(text: str, finish_reason: str | None) -> dict[str, JsonValue]:
|
||||
return {"text": text, "index": 0, "logprobs": None, "finish_reason": finish_reason}
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
body: Final = JSON_OBJECT.validate_json(request.body)
|
||||
assert body.get("stream") is True, body
|
||||
assert object_value(body["stream_options"]).get("include_usage") is True, body
|
||||
identity: Final = body["prompt"]
|
||||
assert isinstance(identity, str), body
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
*(_text_frame(identity, [_text_choice(piece, None)]) for piece in _PIECES),
|
||||
_text_frame(identity, [_text_choice("", "length")]),
|
||||
_text_frame(identity, [], usage=_USAGE),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _config(directory: Path, wire: Wire) -> Path:
|
||||
base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
|
||||
config: Final = {
|
||||
**base,
|
||||
"model_list": [
|
||||
*list_value(base["model_list"]),
|
||||
{
|
||||
"model_name": _TEXT_MODEL,
|
||||
"litellm_params": {
|
||||
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
|
||||
"api_base": f"{wire.url}/v1",
|
||||
"api_key": _PROVIDER_KEY,
|
||||
},
|
||||
},
|
||||
],
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": _GUARDRAIL,
|
||||
"litellm_params": {
|
||||
"guardrail": "tool_permission",
|
||||
"mode": "post_call",
|
||||
"default_on": True,
|
||||
"default_action": "allow",
|
||||
"rules": [],
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
path: Final = directory / f"text-rebuild-{uuid.uuid4().hex}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("text-rebuild-rig")
|
||||
with gateway_from_environment() as gateway, wire_server(_respond) as wire:
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
directory,
|
||||
{},
|
||||
config=_config(directory, wire),
|
||||
workers=2,
|
||||
healthcheck=WorkerHealthcheck.LENIENT,
|
||||
client_timeout=_CELL_SECONDS,
|
||||
) as owned:
|
||||
yield _Rig(wire, owned)
|
||||
|
||||
|
||||
def _stream(rig: _Rig, identity: str) -> _Stream:
|
||||
gateway: Final = rig.proxy.gateway
|
||||
with gateway.client.stream(
|
||||
"POST",
|
||||
"/v1/completions",
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
json={
|
||||
"model": _TEXT_MODEL,
|
||||
"prompt": identity,
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
) as response:
|
||||
return _Stream(response.status_code, response.read())
|
||||
|
||||
|
||||
def _posted_targets(wire: Wire) -> tuple[str, ...]:
|
||||
return tuple(request.target for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
def test_guarded_text_completion_stream_reaches_the_caller_whole_with_its_usage(rig: _Rig) -> None:
|
||||
rig.wire.drain()
|
||||
identity: Final = f"cmpl-guarded-{uuid.uuid4().hex}"
|
||||
|
||||
stream: Final = _stream(rig, identity)
|
||||
|
||||
assert stream.status == 200, stream.body
|
||||
assert stream.body.rstrip().endswith(b"data: [DONE]"), stream.body[-400:]
|
||||
assert stream.ids == {identity}, stream.ids
|
||||
assert stream.text == "".join(_PIECES), stream.frames
|
||||
assert stream.finish_reasons == ("length",), stream.frames
|
||||
assert [(usage["prompt_tokens"], usage["completion_tokens"]) for usage in stream.usages] == [
|
||||
(_PROMPT_TOKENS, _COMPLETION_TOKENS)
|
||||
], stream.frames
|
||||
assert _posted_targets(rig.wire) == ("/v1/completions",)
|
||||
(row,) = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s',
|
||||
(identity,),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _COMPLETION_TOKENS), row
|
||||
assert _GUARDRAIL in list_value(object_value(row["metadata"])["applied_guardrails"]), row
|
||||
|
||||
|
||||
def _together_chunk(choices: list[JsonValue], **extra: JsonValue) -> bytes:
|
||||
return _frame(
|
||||
{
|
||||
"id": f"chatcmpl-together-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "scripted",
|
||||
"choices": choices,
|
||||
**extra,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _together_chat_stream(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target.endswith("/chat/completions"), request
|
||||
assert JSON_OBJECT.validate_json(request.body).get("stream") is True, request.body
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
_together_chunk(
|
||||
[
|
||||
{
|
||||
"index": 0,
|
||||
"text": _BRIDGE_TEXT,
|
||||
"delta": {"role": "assistant", "content": _BRIDGE_TEXT},
|
||||
"finish_reason": None,
|
||||
}
|
||||
]
|
||||
),
|
||||
_together_chunk([{"index": 0, "text": "", "delta": {}, "finish_reason": "stop"}]),
|
||||
_together_chunk([], usage=_USAGE),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _completed_events(body: bytes) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
object_value(event["response"])
|
||||
for event in _Stream(200, body).frames
|
||||
if event.get("type") == "response.completed"
|
||||
)
|
||||
|
||||
|
||||
def _output_text(response: Mapping[str, JsonValue]) -> str:
|
||||
parts: Final = chain.from_iterable(
|
||||
_content_parts(object_value(item)) for item in list_value(response.get("output") or [])
|
||||
)
|
||||
return "".join(str(part.get("text") or "") for part in parts)
|
||||
|
||||
|
||||
def _content_parts(item: Mapping[str, JsonValue]) -> Iterator[Mapping[str, JsonValue]]:
|
||||
return (object_value(part) for part in list_value(item.get("content") or []))
|
||||
|
||||
|
||||
def test_responses_bridge_rebuilds_a_chat_stream_whose_choices_also_carry_text(gateway: Gateway) -> None:
|
||||
with wire_server(_together_chat_stream) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_TOGETHER_MODEL, api_base=f"{wire.url}/v1", api_key=_PROVIDER_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/responses",
|
||||
json={"model": model, "input": f"bridge probe {uuid.uuid4().hex}", "stream": True},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
(completed,) = _completed_events(response.content)
|
||||
assert _output_text(completed) == _BRIDGE_TEXT, completed
|
||||
usage: Final = object_value(completed["usage"])
|
||||
assert (usage["input_tokens"], usage["output_tokens"]) == (_PROMPT_TOKENS, _COMPLETION_TOKENS), usage
|
||||
assert [request.target for request in wire.drain()] == ["/v1/chat/completions"]
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
|
@ -12,7 +12,7 @@ import respx
|
|||
import litellm
|
||||
from litellm import ChatCompletionUsageBlock, stream_chunk_builder
|
||||
from litellm.types.utils import GenericStreamingChunk
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor, stream_chunk_builder_text_completion
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionDeltaToolCall,
|
||||
|
|
@ -23,6 +23,8 @@ from litellm.types.utils import (
|
|||
PromptTokensDetails,
|
||||
ServerToolUse,
|
||||
StreamingChoices,
|
||||
TextChoices,
|
||||
TextCompletionResponse,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
|
@ -2140,6 +2142,341 @@ def test_stream_chunk_builder_joins_openai_audio_output_pieces() -> None:
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reported_total", [{"total_tokens": 13}, {}])
|
||||
@pytest.mark.parametrize("trailer_choices", [[], [{"text": None, "index": 0, "logprobs": None, "finish_reason": None}]])
|
||||
def test_stream_chunk_builder_text_completion_keeps_finish_reason_and_provider_usage(
|
||||
trailer_choices: list[dict[str, object]], reported_total: dict[str, int]
|
||||
) -> None:
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder_text_completion
|
||||
from litellm.types.utils import TextCompletionResponse
|
||||
|
||||
def chunk(choices: list[dict[str, object]], **extra: object) -> TextCompletionResponse:
|
||||
return TextCompletionResponse(
|
||||
id="cmpl-1", object="text_completion", created=1, model="my-model", choices=choices, **extra
|
||||
)
|
||||
|
||||
chunks: Final = [
|
||||
chunk([{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}]),
|
||||
chunk([{"text": " world", "index": 0, "logprobs": None, "finish_reason": None}]),
|
||||
chunk([{"text": "", "index": 0, "logprobs": None, "finish_reason": "length"}]),
|
||||
chunk(
|
||||
trailer_choices,
|
||||
usage={
|
||||
"prompt_tokens": 7,
|
||||
"completion_tokens": 6,
|
||||
**reported_total,
|
||||
"prompt_tokens_details": {"cached_tokens": 4},
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
response: Final = stream_chunk_builder_text_completion(chunks=chunks, messages=None)
|
||||
|
||||
assert response.choices[0].text == "Hello world"
|
||||
assert response.choices[0].finish_reason == "length"
|
||||
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (7, 6, 13)
|
||||
assert response.usage.prompt_tokens_details.cached_tokens == 4
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_usage_dict_set_on_a_typed_text_completion_chunk() -> None:
|
||||
def chunk(choices: list[dict[str, object]]) -> TextCompletionResponse:
|
||||
return TextCompletionResponse(
|
||||
id="cmpl-1", object="text_completion", created=1, model="my-model", choices=choices
|
||||
)
|
||||
|
||||
trailer: Final = chunk([{"text": "", "index": 0, "logprobs": None, "finish_reason": "length"}])
|
||||
trailer["usage"] = {"prompt_tokens": 7, "completion_tokens": 6, "total_tokens": 13}
|
||||
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[chunk([{"text": "Hello world", "index": 0, "logprobs": None, "finish_reason": None}]), trailer]
|
||||
)
|
||||
|
||||
assert response.choices[0].finish_reason == "length"
|
||||
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (7, 6, 13)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("build", [litellm.stream_chunk_builder, litellm.stream_chunk_builder_text_completion])
|
||||
def test_stream_chunk_builder_rebuilds_dict_text_completion_chunks(build: Callable[..., TextCompletionResponse]) -> None:
|
||||
def chunk(choices: list[dict[str, object]], **extra: object) -> dict[str, object]:
|
||||
return {
|
||||
"id": "cmpl-1",
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": "my-model",
|
||||
"choices": choices,
|
||||
**extra,
|
||||
}
|
||||
|
||||
chunks: Final = [
|
||||
chunk([{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}]),
|
||||
chunk([{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}]),
|
||||
chunk([], usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}),
|
||||
]
|
||||
|
||||
response: Final = build(chunks=chunks)
|
||||
|
||||
assert response.id == "cmpl-1"
|
||||
assert response.choices[0].text == "Hello world"
|
||||
assert response.choices[0].finish_reason == "stop"
|
||||
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (3, 2, 5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("choice_text", ["Hi", None])
|
||||
def test_stream_chunk_builder_keeps_dumped_chat_chunks_with_a_choice_text_on_the_chat_path(
|
||||
choice_text: str | None,
|
||||
) -> None:
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
def dumped_chunk(content: str, finish_reason: str | None) -> dict[str, object]:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-1",
|
||||
object="chat.completion.chunk",
|
||||
created=1,
|
||||
model="my-model",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"text": choice_text,
|
||||
"delta": {"role": "assistant", "content": content},
|
||||
"finish_reason": finish_reason,
|
||||
"logprobs": None,
|
||||
}
|
||||
],
|
||||
).model_dump()
|
||||
|
||||
response: Final = litellm.stream_chunk_builder(chunks=[dumped_chunk("Hi", None), dumped_chunk(" there", "stop")])
|
||||
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
assert response.choices[0].message.content == "Hi there"
|
||||
assert response.choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_logprobs_from_every_text_completion_chunk() -> None:
|
||||
def chunk(text: str, offset: int, logprob: float, finish_reason: str | None) -> dict[str, object]:
|
||||
return {
|
||||
"id": "cmpl-1",
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": "my-model",
|
||||
"choices": [
|
||||
{
|
||||
"text": text,
|
||||
"index": 0,
|
||||
"logprobs": {
|
||||
"tokens": [text],
|
||||
"token_logprobs": [logprob],
|
||||
"top_logprobs": [{text: logprob}],
|
||||
"text_offset": [offset],
|
||||
},
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
trailer: Final = {
|
||||
"id": "cmpl-1",
|
||||
"object": "text_completion",
|
||||
"created": 1,
|
||||
"model": "my-model",
|
||||
"choices": [],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3},
|
||||
}
|
||||
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[chunk("Hello", 5, -0.1, None), chunk(" world", 10, -0.2, "length"), trailer]
|
||||
)
|
||||
|
||||
assert response.choices[0].text == "Hello world"
|
||||
assert response.choices[0].logprobs.model_dump() == {
|
||||
"tokens": ["Hello", " world"],
|
||||
"token_logprobs": [-0.1, -0.2],
|
||||
"top_logprobs": [{"Hello": -0.1}, {" world": -0.2}],
|
||||
"text_offset": [5, 10],
|
||||
}
|
||||
|
||||
|
||||
def _text_completion_chunk(*choices: dict[str, object]) -> dict[str, object]:
|
||||
return {"id": "cmpl-1", "object": "text_completion", "created": 1, "model": "my-model", "choices": list(choices)}
|
||||
|
||||
|
||||
def _text_choice(
|
||||
index: int, text: str, offset: int, top_logprobs: list[dict[str, float]] | None, finish_reason: str | None = None
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"text": text,
|
||||
"index": index,
|
||||
"logprobs": {"tokens": [text], "token_logprobs": [-0.5], "top_logprobs": top_logprobs, "text_offset": [offset]},
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
|
||||
|
||||
def test_stream_chunk_builder_rebuilds_each_text_completion_choice_by_index() -> None:
|
||||
usage_trailer: Final = {
|
||||
**_text_completion_chunk(),
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 4, "total_tokens": 5},
|
||||
}
|
||||
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[
|
||||
_text_completion_chunk(_text_choice(0, "Hi", 5, [{"Hi": -0.5}])),
|
||||
_text_completion_chunk(_text_choice(1, "Hey", 5, [{"Hey": -0.5}])),
|
||||
_text_completion_chunk(_text_choice(0, " there", 7, [{" there": -0.5}], "stop")),
|
||||
_text_completion_chunk(_text_choice(1, " you", 8, [{" you": -0.5}], "length")),
|
||||
usage_trailer,
|
||||
]
|
||||
)
|
||||
|
||||
assert [
|
||||
(choice.index, choice.text, choice.finish_reason, choice.logprobs.tokens, choice.logprobs.text_offset)
|
||||
for choice in response.choices
|
||||
] == [(0, "Hi there", "stop", ["Hi", " there"], [5, 7]), (1, "Hey you", "length", ["Hey", " you"], [5, 8])]
|
||||
assert (response.usage.prompt_tokens, response.usage.completion_tokens) == (1, 4)
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_null_text_completion_logprobs_fields_null() -> None:
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[
|
||||
_text_completion_chunk(_text_choice(0, "Hello", 5, None)),
|
||||
_text_completion_chunk(_text_choice(0, " world", 10, None, "length")),
|
||||
]
|
||||
)
|
||||
|
||||
assert response.choices[0].logprobs.model_dump() == {
|
||||
"tokens": ["Hello", " world"],
|
||||
"token_logprobs": [-0.5, -0.5],
|
||||
"top_logprobs": None,
|
||||
"text_offset": [5, 10],
|
||||
}
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_a_text_completion_stream_that_ends_on_a_chat_chunk() -> None:
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[
|
||||
TextCompletionResponse(id="cmpl-1", model="my-model", choices=[TextChoices(index=0, text="Hello")]),
|
||||
ModelResponseStream(
|
||||
id="cmpl-1",
|
||||
model="my-model",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")],
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert isinstance(response, TextCompletionResponse)
|
||||
assert [(choice.text, choice.finish_reason) for choice in response.choices] == [("Hello", "stop")]
|
||||
|
||||
|
||||
def test_stream_chunk_builder_drops_only_the_text_completion_logprobs_it_cannot_read() -> None:
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[
|
||||
_text_completion_chunk({"text": "Hello", "index": 0, "logprobs": {"tokens": ["Hello"]}}),
|
||||
_text_completion_chunk(_text_choice(0, " world", 10, [{" world": -0.5}], "stop")),
|
||||
]
|
||||
)
|
||||
|
||||
assert response.choices[0].text == "Hello world"
|
||||
assert response.choices[0].logprobs.tokens == [" world"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"first_logprobs",
|
||||
[[-0.1], {"tokens": ["Hello"], "token_logprobs": [-0.1], "top_logprobs": None, "text_offset": [0]}],
|
||||
ids=["list", "logprobs_object"],
|
||||
)
|
||||
def test_stream_chunk_builder_keeps_text_completion_logprobs_of_another_shape(
|
||||
first_logprobs: Sequence[float] | Mapping[str, object],
|
||||
) -> None:
|
||||
def chunk(
|
||||
text: str, logprobs: Sequence[float] | Mapping[str, object], finish_reason: str | None
|
||||
) -> TextCompletionResponse:
|
||||
return TextCompletionResponse(
|
||||
id="cmpl-1",
|
||||
model="my-model",
|
||||
choices=[TextChoices(index=0, text=text, logprobs=logprobs, finish_reason=finish_reason)],
|
||||
)
|
||||
|
||||
response: Final = litellm.stream_chunk_builder(
|
||||
chunks=[chunk("Hello", first_logprobs, None), chunk(" world", [-0.2], "stop")]
|
||||
)
|
||||
|
||||
assert (response.choices[0].text, response.choices[0].logprobs) == ("Hello world", [-0.2])
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_the_usage_a_chat_chunk_reports_in_a_text_completion_stream() -> None:
|
||||
response: Final = stream_chunk_builder_text_completion(
|
||||
chunks=[
|
||||
TextCompletionResponse(
|
||||
id="cmpl-1", model="my-model", choices=[TextChoices(index=0, text="Hello", finish_reason="stop")]
|
||||
),
|
||||
ModelResponseStream(
|
||||
id="cmpl-1",
|
||||
model="my-model",
|
||||
choices=[],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert response.usage is not None
|
||||
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (3, 1, 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first_index", "second_index", "expected"),
|
||||
[
|
||||
(None, None, ({"index": 0, "text": "Hello world", "finish_reason": "stop"},)),
|
||||
("0", 0, ({"index": 0, "text": "Hello world", "finish_reason": "stop"},)),
|
||||
("0", "0", ({"index": 0, "text": "Hello world", "finish_reason": "stop"},)),
|
||||
(float("inf"), 0, ({"index": 0, "text": "Hello world", "finish_reason": "stop"},)),
|
||||
(
|
||||
0,
|
||||
"1",
|
||||
(
|
||||
{"index": 0, "text": "Hello", "finish_reason": None},
|
||||
{"index": 1, "text": " world", "finish_reason": "stop"},
|
||||
),
|
||||
),
|
||||
],
|
||||
ids=["none", "string_then_int", "strings", "infinity_then_int", "int_then_string"],
|
||||
)
|
||||
def test_stream_chunk_builder_rebuilds_text_completion_choice_indexes_that_are_not_ints_as_ints(
|
||||
first_index: float | str | None,
|
||||
second_index: int | str | None,
|
||||
expected: tuple[Mapping[str, object], ...],
|
||||
) -> None:
|
||||
response: Final = stream_chunk_builder_text_completion(
|
||||
chunks=[
|
||||
_text_completion_chunk({"text": "Hello", "index": first_index}),
|
||||
_text_completion_chunk({"text": " world", "index": second_index, "finish_reason": "stop"}),
|
||||
]
|
||||
)
|
||||
|
||||
assert tuple(choice.model_dump(include={"index", "text", "finish_reason"}) for choice in response.choices) == expected
|
||||
|
||||
|
||||
def test_stream_chunk_builder_text_completion_survives_uncountable_prompt() -> None:
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder_text_completion
|
||||
from litellm.types.utils import TextCompletionResponse
|
||||
|
||||
chunks: Final = [
|
||||
TextCompletionResponse(
|
||||
id="cmpl-1",
|
||||
object="text_completion",
|
||||
created=1,
|
||||
model="my-model",
|
||||
choices=[{"text": "Hello world", "index": 0, "logprobs": None, "finish_reason": "stop"}],
|
||||
)
|
||||
]
|
||||
|
||||
response: Final = stream_chunk_builder_text_completion(
|
||||
chunks=chunks, messages=[{"role": "user", "content": [{"type": "not-a-content-type"}]}]
|
||||
)
|
||||
|
||||
assert response.choices[0].text == "Hello world"
|
||||
assert response.usage.prompt_tokens == 0
|
||||
assert response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens == response.usage.completion_tokens
|
||||
|
||||
|
||||
def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies():
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
"""Unit tests for the Microsoft Purview DLP guardrail."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -536,6 +538,61 @@ class TestTextCompletionHooks:
|
|||
assert "alpha" in combined
|
||||
assert "beta" in combined
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_multi_choice_text_completion_is_blocked_like_the_non_streamed_one(self):
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.types.utils import TextChoices, TextCompletionResponse
|
||||
|
||||
sensitive: Final = "4111 1111 1111 1111"
|
||||
choice_pieces: Final = ((sensitive[:9], sensitive[9:]), ("Hey", " you"))
|
||||
user: Final = UserAPIKeyAuth(api_key="test", user_id="user-123")
|
||||
|
||||
def graph(request: httpx.Request) -> httpx.Response:
|
||||
if request.url.host == "login.microsoftonline.com":
|
||||
return httpx.Response(200, json={"access_token": "graph-token", "expires_in": 3600})
|
||||
if not request.url.path.endswith("/processContent"):
|
||||
return httpx.Response(200, json={"value": []}, headers={"ETag": "scope-etag"})
|
||||
entry: Final = json.loads(request.content)["contentToProcess"]["contentEntries"][0]
|
||||
block: Final = {"@odata.type": "#microsoft.graph.restrictAccessAction", "restrictionAction": "block"}
|
||||
actions: Final = [block] if sensitive in entry["content"]["data"] else []
|
||||
return httpx.Response(200, json={"protectionScopeState": "unchanged", "policyActions": actions})
|
||||
|
||||
def guardrail_on_graph() -> MicrosoftPurviewDLPGuardrail:
|
||||
guardrail: Final = _make_guardrail()
|
||||
guardrail.async_handler = AsyncHTTPHandler(transport=httpx.MockTransport(graph))
|
||||
return guardrail
|
||||
|
||||
def chunk(index: int, text: str) -> TextCompletionResponse:
|
||||
return TextCompletionResponse(
|
||||
id="cmpl-n2", model="gpt-3.5-turbo-instruct", choices=[TextChoices(text=text, index=index)]
|
||||
)
|
||||
|
||||
async def interleaved_stream():
|
||||
for step in range(2):
|
||||
for index, pieces in enumerate(choice_pieces):
|
||||
yield chunk(index, pieces[step])
|
||||
|
||||
non_streamed: Final = TextCompletionResponse(
|
||||
id="cmpl-n2",
|
||||
model="gpt-3.5-turbo-instruct",
|
||||
choices=[TextChoices(text="".join(pieces), index=index) for index, pieces in enumerate(choice_pieces)],
|
||||
)
|
||||
with pytest.raises(HTTPException) as non_streamed_block:
|
||||
await guardrail_on_graph().async_post_call_success_hook(
|
||||
data={}, user_api_key_dict=user, response=non_streamed
|
||||
)
|
||||
|
||||
async def drain_guarded_stream() -> None:
|
||||
async for _ in guardrail_on_graph().async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user, response=interleaved_stream(), request_data={}
|
||||
):
|
||||
pass
|
||||
|
||||
with pytest.raises(HTTPException) as streamed_block:
|
||||
await drain_guarded_stream()
|
||||
|
||||
assert streamed_block.value.status_code == non_streamed_block.value.status_code == 400
|
||||
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# Responses API hooks
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue