diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d505192184e..b42210e2793 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -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"] diff --git a/litellm/main.py b/litellm/main.py index 8e9ff3cf910..65c504072a9 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 ########## diff --git a/tests/integration/sdk/test_text_completion_stream_rebuild_sdk.py b/tests/integration/sdk/test_text_completion_stream_rebuild_sdk.py new file mode 100644 index 00000000000..9ece946abca --- /dev/null +++ b/tests/integration/sdk/test_text_completion_stream_rebuild_sdk.py @@ -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 diff --git a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py new file mode 100644 index 00000000000..60a33d46b29 --- /dev/null +++ b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py @@ -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"] diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index f8cd0416b1c..5c56245cf33 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index 4a7a14fceaa..25d2556bf26 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -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