scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge 8807b8bfe3 into 0b39e5faf3

Browse source
This commit is contained in:
AlbertNeuralTrust 2026-11-03 17:20:14 -08:00 • committed by GitHub
commit d2cdc5c631
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1388 additions and 288 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/integrations/custom_guardrail.py
View file

@ -1533,6 +1533,7 @@ class CustomGuardrail(CustomLogger):
end_time=end_time,
duration=duration,
masked_entity_count=masked_entity_count,
choice_index=request_data.get("guardrail_choice_index"),
**(tracing_detail or {}),
)

1
litellm/integrations/otel/mappers/genai.py
View file

@ -141,6 +141,7 @@ class GenAIMapper:
LiteLLM.GUARDRAIL_CONFIDENCE_SCORE: lambda d: d.confidence_score,
LiteLLM.GUARDRAIL_RISK_SCORE: lambda d: d.risk_score,
LiteLLM.GUARDRAIL_MASKED_ENTITY_COUNT: lambda d: d.masked_entity_count,
LiteLLM.GUARDRAIL_CHOICE_INDEX: lambda d: d.choice_index,
LiteLLM.GUARDRAIL_DURATION: lambda d: d.duration,
LiteLLM.GUARDRAIL_ID: lambda d: d.guardrail_id,
LiteLLM.GUARDRAIL_POLICY_TEMPLATE: lambda d: d.policy_template,

2
litellm/integrations/otel/model/payloads.py
View file

@ -253,6 +253,7 @@ class GuardrailSpanData:
mode: str | None = None
status: str | None = None
masked_entity_count: int | None = None
choice_index: int | None = None
provider: str | None = None
action: str | None = None
# The guardrail verdict / provider response (e.g. the moderation result),
@ -313,6 +314,7 @@ class GuardrailSpanData:
mode=_guardrail_mode_str(get("guardrail_mode")),
status=status,
masked_entity_count=_total_masked_entities(get("masked_entity_count")),
choice_index=as_int(get("choice_index")),
provider=as_str(get("guardrail_provider")),
action=as_str(get("guardrail_action")),
response_json=_json_or_none(response) if response is not None else None,

1
litellm/integrations/otel/model/semconv.py
View file

@ -317,6 +317,7 @@ class LiteLLM:
GUARDRAIL_CONFIDENCE_SCORE: Final = "litellm.guardrail.confidence_score"
GUARDRAIL_RISK_SCORE: Final = "litellm.guardrail.risk_score"
GUARDRAIL_MASKED_ENTITY_COUNT: Final = "litellm.guardrail.masked_entity_count"
GUARDRAIL_CHOICE_INDEX: Final = "litellm.guardrail.choice_index"
GUARDRAIL_DURATION: Final = "litellm.guardrail.duration"
GUARDRAIL_ID: Final = "litellm.guardrail.id"
GUARDRAIL_POLICY_TEMPLATE: Final = "litellm.guardrail.policy_template"

6
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
@ -34,6 +34,10 @@ class StreamTransformSink:
mutated_text_per_choice: dict[int, str] = field(default_factory=dict)
holdback_per_choice: dict[int, int] = field(default_factory=dict)
def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None:
self.mutated_text_per_choice = dict(texts)
self.holdback_per_choice = dict(holdback)
@dataclass(frozen=True, slots=True)
class StreamingScanKey:

500
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -17,10 +17,13 @@ This pattern can be replicated for other message formats (e.g., Anthropic).
import json
import time
import uuid
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Generator, Iterator, Mapping, Sequence
from contextlib import contextmanager
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union, cast
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
@ -50,7 +53,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_items,
unappliable_request_rewrite,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
coerce_stream_holdback_value,
)
@ -76,6 +79,30 @@ if TYPE_CHECKING:
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_INSPECTION_KEYS: Final[tuple[str, ...]] = ("response", "responses", "guardrail_choice_index")
@contextmanager
def _borrowed_inspection_keys(request_data: dict[str, object]) -> Generator[None]:
"""Lend a round the keys it inspects through, and give the caller's own back.
The round owns ``response``, ``responses`` and the choice label only while it runs. What it
leaves behind is read by later hooks, and a label the caller supplied is not one a round may
report, so the incoming label is dropped for the duration.
"""
missing: Final = object()
priors: Final = tuple((key, request_data.get(key, missing)) for key in _INSPECTION_KEYS)
request_data.pop("guardrail_choice_index", None)
try:
yield
finally:
for key, prior in priors:
if prior is missing:
request_data.pop(key, None)
else:
request_data[key] = prior
class OpenAIChatCompletionsHandler(BaseTranslation):
"""
Handler for processing OpenAI chat completions messages with guardrails.
@ -172,8 +199,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
logging_obj=litellm_logging_obj,
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", [])
guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ())
guardrailed_tools: Final = guardrailed_inputs.get("tools")
if guardrailed_tools is not None:
data["tools"] = (
@ -392,7 +419,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
request_data: dict[str, object] | None = None,
) -> ModelResponse:
"""
Process output response by applying guardrails to text content.
@ -443,9 +470,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# create a standalone dict (SDK / direct-call path).
if request_data is None:
request_data = {"response": response}
else:
if "response" not in request_data:
request_data["response"] = response
request_data.setdefault("response", response)
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
@ -465,7 +490,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
logging_obj=litellm_logging_obj,
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
guardrailed_tool_calls: Final[list[dict[str, object]]] = (
cast(list[dict[str, object]], returned_tool_calls)
@ -499,7 +524,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
request_data: dict[str, object] | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
@ -555,14 +580,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict | None,
request_data: dict[str, object] | None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
terminate the stream. Text rewrites are not propagated to the client here
(see ``_process_streaming_transform`` for the incremental_diff path) unless
``deliver_ended_stream_rewrites`` opts the ended-stream branch in."""
has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far)
if has_stream_ended:
await self._process_ended_stream(
@ -609,8 +634,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# create a standalone dict (SDK / direct-call path).
if request_data is None:
request_data = {"responses": responses_so_far}
owes_responses = False # rebind-ok: the standalone dict is ours to keep
else:
if "responses" not in request_data:
owes_responses = "responses" not in request_data # rebind-ok: tracks the key this scan adds
if owes_responses:
request_data["responses"] = responses_so_far
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
@ -621,14 +648,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# Include model information from the first response if available
if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model:
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
try:
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
finally:
if owes_responses:
request_data.pop("responses", None)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
guardrailed_texts: Final = guardrailed_inputs.get("texts", ())
# Step 4: Apply guardrailed text back to all streaming chunks
# For each choice, replace the combined text across all chunks
@ -661,13 +692,45 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
await self.process_output_response(
response=model_response,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
inspection_responses: Final = (
tuple(
model_response.model_copy(update=MappingProxyType({"choices": [choice]}))
for choice in model_response.choices
)
if pre_guardrail_tool_calls and len(model_response.choices) > 1
else (model_response,)
)
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
inspection_request_data: Final = request_data if request_data is not None else {}
undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure
with _borrowed_inspection_keys(inspection_request_data):
for inspection_response in inspection_responses:
inspection_request_data["response"] = inspection_response
inspection_request_data["guardrail_choice_index"] = (
inspection_response.choices[0].index if len(inspection_response.choices) == 1 else None
)
inspection_request_data["responses"] = (
[
self._narrowed_to_choice(chunk, inspection_response.choices[0].index)
for chunk in responses_so_far
]
if len(inspection_response.choices) == 1
else responses_so_far
)
try:
await self.process_output_response(
response=inspection_response,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=inspection_request_data,
)
except UndeliverableStreamRewrite as exc:
if undeliverable is None:
undeliverable = exc
if undeliverable is not None:
raise undeliverable
if not deliver_ended_stream_rewrites:
return
await self._write_ended_stream_text_rewrites(
@ -682,6 +745,96 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
@staticmethod
def _streamed_argument_fragments(
responses_so_far: Sequence["ModelResponseStream"],
) -> Iterator[tuple[tuple[int, int], str]]:
for chunk in responses_so_far:
for choice in chunk.choices:
for call in choice.delta.tool_calls or ():
arguments = getattr(getattr(call, "function", None), "arguments", None)
yield (choice.index, call.index or 0), arguments if isinstance(arguments, str) else ""
@staticmethod
def _parses_as_json(arguments: str) -> bool:
if not arguments:
return False
try:
json.loads(arguments)
except ValueError:
return False
return True
@staticmethod
def _joined_streamed_arguments(
responses_so_far: Sequence["ModelResponseStream"],
) -> Mapping[tuple[int, int], str]:
"""The argument text streamed so far per (choice index, tool call index).
Grouped in one pass. Rescanning every fragment once per distinct call made a sampled
round cost the number of calls times the length of the stream, which at n=2 grows with
the square of the streamed arguments.
"""
fragments: Final = sorted(
OpenAIChatCompletionsHandler._streamed_argument_fragments(responses_so_far),
key=lambda fragment: fragment[0],
)
return MappingProxyType(
{call: "".join(text for _, text in group) for call, group in groupby(fragments, key=lambda f: f[0])}
)
@staticmethod
def _streamed_tool_arguments_are_complete(responses_so_far: Sequence["ModelResponseStream"]) -> bool:
"""Whether every streamed tool call now carries arguments that parse as JSON.
A call that has announced itself but emitted no argument text yet is still in flight, so a
parallel call finishing first does not make the whole set complete.
"""
joined: Final = OpenAIChatCompletionsHandler._joined_streamed_arguments(responses_so_far)
if not joined:
return False
return all(OpenAIChatCompletionsHandler._parses_as_json(arguments) for arguments in joined.values())
@staticmethod
def _streamed_call_ids(responses_so_far: Sequence["ModelResponseStream"]) -> Mapping[tuple[int, int], str]:
"""The id each (choice index, tool call index) announced itself with."""
choices: Final = chain.from_iterable(chunk.choices for chunk in responses_so_far)
calls: Final = chain.from_iterable(
((choice.index, call) for call in choice.delta.tool_calls or ()) for choice in choices
)
return MappingProxyType({(index, call.index or 0): call.id for index, call in calls if call.id})
@staticmethod
def _choice_without(choice: "Choices", unfinished: frozenset[tuple[int, str]]) -> "Choices":
kept: Final = [call for call in choice.message.tool_calls or () if (choice.index, call.id) not in unfinished]
message: Final = choice.message.model_copy(update=MappingProxyType({"tool_calls": kept or None}))
return choice.model_copy(update=MappingProxyType({"message": message}))
@staticmethod
def _without_unfinished_tool_calls(
assembled: "ModelResponse", responses_so_far: Sequence["ModelResponseStream"]
) -> "ModelResponse":
"""Copy of the rebuilt response without the calls still in flight when the stream ended.
Rebuilding gives a call that never streamed complete arguments an empty JSON object, which
reads to a guardrail as a real call taking no arguments rather than an unfinished one. The
calls are matched by the choice they belong to and the id they announced, since the rebuilt
order need not follow the order the stream gave them, and providers number tool calls per
choice, so the same id on another choice is another call.
"""
joined: Final = OpenAIChatCompletionsHandler._joined_streamed_arguments(responses_so_far)
unfinished: Final = frozenset(
(key[0], call_id)
for key, call_id in OpenAIChatCompletionsHandler._streamed_call_ids(responses_so_far).items()
if not OpenAIChatCompletionsHandler._parses_as_json(joined.get(key, ""))
)
if not unfinished:
return assembled
pruned: Final = [
OpenAIChatCompletionsHandler._choice_without(choice, unfinished) for choice in assembled.choices
]
return assembled.model_copy(update=MappingProxyType({"choices": pruned}))
@staticmethod
def _rebuild_ended_stream_per_choice(
responses_so_far: Sequence["ModelResponseStream"],
@ -706,6 +859,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
)
for index in choice_indices
)
if not rebuilt_by_index:
whole: Final = stream_chunk_builder(chunks=list(responses_so_far), logging_obj=litellm_logging_obj)
return cast(ModelResponse, whole) if whole is not None else ModelResponse(choices=[])
(_, base_response), *_ = rebuilt_by_index
stitched_choices: Final = [
rebuilt.choices[0].model_copy(update=MappingProxyType({"index": index}))
@ -763,7 +919,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict | None,
request_data: dict[str, object] | None,
sink: StreamTransformSink,
) -> None:
"""Run the guardrail over the raw accumulated text and report the
@ -775,8 +931,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""
raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far)
if not raw_by_index:
sink.mutated_text_per_choice = {}
sink.holdback_per_choice = {}
sink.record(MappingProxyType({}), MappingProxyType({}))
return
# Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits
@ -788,54 +943,145 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
indices: Final = sorted(raw_by_index.keys())
texts_to_check: Final = [raw_by_index[i] for i in indices]
if request_data is None:
request_data = {"responses": responses_so_far}
elif "responses" not in request_data:
request_data["responses"] = responses_so_far
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if responses_so_far and getattr(responses_so_far[0], "model", None):
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
resolved_request_data: Final[dict[str, object]] = (
request_data if request_data is not None else {"responses": responses_so_far}
)
with _borrowed_inspection_keys(resolved_request_data):
if len(indices) == 1:
resolved_request_data["guardrail_choice_index"] = indices[0]
if "responses" not in resolved_request_data:
resolved_request_data["responses"] = responses_so_far
self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if self._stream_carries_tool_calls(responses_so_far):
stream_ended: Final = self._first_choice_has_finished(responses_so_far)
tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far)
assembled: Final = (
None
if not stream_ended and not tool_arguments_complete
else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
)
inspected: Final = (
None
if assembled is None
else assembled
if tool_arguments_complete
else self._without_unfinished_tool_calls(assembled, responses_so_far)
)
if inspected is not None:
resolved_request_data["response"] = inspected
if inspected is not None and len(inspected.choices) > 1:
choice_rounds: Final = tuple(
(
choice,
StreamTransformSink(),
[self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far],
)
for choice in inspected.choices
)
from litellm.proxy.policy_engine.pipeline_executor import (
UndeliverableStreamRewrite,
)
returned_texts = guardrailed_inputs.get("texts")
# No "texts" key means the guardrail made no change (action NONE): the raw
# accumulated text is the guardrailed text. A present-but-shorter list is a
# guardrail contract violation; those choices are omitted below (withheld,
# not emitted raw) so a malformed response fails closed instead of leaking.
if returned_texts is None:
returned_texts = texts_to_check
elif len(returned_texts) < len(texts_to_check):
verbose_proxy_logger.warning(
"OpenAI Chat Completions: guardrail returned %s transformed texts for %s inputs on the "
"streaming transform path; withholding the unmatched choices to fail closed.",
len(returned_texts),
len(texts_to_check),
transform_undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: first failure
try:
for choice, choice_sink, choice_chunks in choice_rounds:
resolved_request_data["response"] = inspected.model_copy(
update=MappingProxyType({"choices": [choice]})
)
resolved_request_data["responses"] = choice_chunks
resolved_request_data["guardrail_choice_index"] = choice.index
try:
await self._process_streaming_transform(
responses_so_far=choice_chunks,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=resolved_request_data,
sink=choice_sink,
)
except UndeliverableStreamRewrite as exc:
if transform_undeliverable is None:
transform_undeliverable = exc
finally:
resolved_request_data["response"] = assembled
resolved_request_data["responses"] = responses_so_far
resolved_request_data.pop("guardrail_choice_index", None)
if transform_undeliverable is not None:
raise transform_undeliverable
merged_texts: Final = MappingProxyType(
dict(
chain.from_iterable(
choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds
)
)
)
merged_holdback: Final = MappingProxyType(
dict(
chain.from_iterable(
choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds
)
)
)
sink.record(merged_texts, merged_holdback)
return
if assembled is not None and tool_arguments_complete:
tool_calls: Final = chain.from_iterable(
choice.message.tool_calls or () for choice in assembled.choices
)
inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python(
tuple(
MappingProxyType({"index": index, **converted})
for index, tool_call in enumerate(tool_calls)
if (converted := self._convert_tool_call_to_dict(tool_call)) is not None
)
)
if responses_so_far and getattr(responses_so_far[0], "model", None):
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=resolved_request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or []
sink.mutated_text_per_choice = {
idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)
}
sink.holdback_per_choice = {
indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
}
returned_texts = guardrailed_inputs.get("texts")
# No "texts" key means the guardrail made no change (action NONE): the raw
# accumulated text is the guardrailed text. A present-but-shorter list is a
# guardrail contract violation; those choices are omitted below (withheld,
# not emitted raw) so a malformed response fails closed instead of leaking.
if returned_texts is None:
returned_texts = texts_to_check
elif len(returned_texts) < len(texts_to_check):
verbose_proxy_logger.warning(
"OpenAI Chat Completions: guardrail returned %s transformed texts for %s inputs on the "
"streaming transform path; withholding the unmatched choices to fail closed.",
len(returned_texts),
len(texts_to_check),
)
holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or ()
transformed_texts: Final = MappingProxyType(
{idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)}
)
transformed_holdback: Final = MappingProxyType(
{
indices[i]: coerce_stream_holdback_value(holdback[i])
for i in range(len(indices))
if i < len(holdback)
}
)
sink.record(transformed_texts, transformed_holdback)
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
stream_ended: Final = self._first_choice_has_finished(responses_so_far)
tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far)
carries_tool_calls: Final = self._stream_carries_tool_calls(responses_so_far)
return StreamingScanKey(
texts=tuple(self._combine_streaming_texts(chunks).values()),
tool_calls=tool_call_fingerprints if stream_ended else (),
tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
stream_ended=stream_ended,
tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended,
tool_calls_in_flight=carries_tool_calls and not stream_ended,
)
def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]:
@ -859,6 +1105,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta"))
)
@staticmethod
def _stream_carries_tool_calls(responses_so_far: Sequence[object]) -> bool:
"""Whether any chunk so far carried a tool call, without fingerprinting them.
A round only needs to know that tool calls are present; fingerprinting every fragment
dumps and serializes each one, and repeating that as the stream grows is what made a
long tool argument cost far more CPU than the arguments themselves.
"""
return any(
_streamed_delta_tool_calls(stream_item_field(choice, "delta"))
for choice in chain.from_iterable(_stream_chunk_choices(chunk) for chunk in responses_so_far)
)
@staticmethod
def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool:
first_choices: Final = tuple(
@ -1142,20 +1401,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
def _function_tool_call_fragments(
responses_so_far: Sequence["ModelResponseStream"],
) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
"""Group the stream's function tool-call fragments by their tool-call index, in
the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
only the indices the builder keeps (an id and a name somewhere in the stream)."""
fragments: Final = tuple(
tool_call
(choice.index, tool_call)
for response in responses_so_far
for choice in response.choices
for tool_call in choice.delta.tool_calls or ()
if isinstance(tool_call, ChatCompletionDeltaToolCall)
)
identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id)
named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name)
return tuple(
tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key)
for key in sorted(identified & named)
)
def _write_ended_stream_tool_call_rewrites(
@ -1165,28 +1422,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail tool-call rewrites back across the buffered
chunks: the rewritten name and full arguments land in the tool call's first
fragment and the arguments of its later fragments are blanked, mirroring the
text write-back. A rewrite on a stream carrying more than one distinct choice
index, or whose fragments do not line up with the rebuilt tool calls, is
reported as undeliverable, so the pipeline executor discards it and releases
the original chunks."""
post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
if len(stream_choice_indices) != 1:
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(
guardrail_name,
f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written "
"back on single-choice streams",
)
if len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
@ -1195,21 +1434,41 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried "
f"{len(fragments_by_tool_call)}",
)
for before, (name, arguments), fragments in zip(
pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
):
if (name, arguments) == before:
continue
head, *tail = fragments
rewritten: Final = tuple(
(fragments, shape)
for before, shape, fragments in zip(
pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
)
if shape != before
)
for fragments, (name, arguments) in rewritten:
head, last = fragments[0], fragments[-1]
head.function.name = name
head.function.arguments = arguments
for fragment in tail:
fragment.function.arguments = ""
if head is not last:
head.function.arguments = ""
last.function.arguments = arguments
superseded: Final = chain.from_iterable(fragments[1:-1] for fragments, _ in rewritten)
self._drop_tool_call_fragments(responses_so_far, frozenset(id(fragment) for fragment in superseded))
@staticmethod
def _drop_tool_call_fragments(
responses_so_far: list["ModelResponseStream"], # mutable-ok: strips the caller's buffered chunks in place
superseded: frozenset[int],
) -> None:
"""Remove fragments a rewrite superseded, so the stream carries no empty-argument tool deltas."""
if not superseded:
return
for response in responses_so_far:
for choice in response.choices:
present = choice.delta.tool_calls or ()
kept = [call for call in present if id(call) not in superseded]
if len(kept) != len(present):
choice.delta.tool_calls = kept or None
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],
guardrailed_texts: list[str],
guardrailed_texts: Sequence[str],
task_mappings: list[tuple[int, int | None]],
) -> None:
"""
@ -1227,12 +1486,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Override this method to customize how responses are applied to streaming responses.
"""
# Build a mapping of what guardrailed text to use for each (choice_idx, content_idx)
guardrail_map: Final[dict[tuple[int, int | None], str]] = {}
for task_idx, guardrail_response in enumerate(guardrailed_texts):
mapping = task_mappings[task_idx]
choice_idx = cast(int, mapping[0])
content_idx_optional = cast(int | None, mapping[1])
guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response
guardrail_map: Final = MappingProxyType(
{task_mappings[index]: text for index, text in enumerate(guardrailed_texts)}
)
# Track which choices we've already set the guardrailed text for
# Key: (choice_idx, content_idx), Value: boolean (True if already set)
@ -1242,34 +1498,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# choice by its index field: on n>1 streams a chunk usually carries one
# choice at list position 0 whose index names the logical choice.
for response in responses:
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
elif isinstance(choice, litellm.Choices):
content = choice.message.content
else:
continue
for choice, message in (
(choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message)
for choice in response.choices
if isinstance(choice, (StreamingChoices, Choices))
):
content = message.content
if content is None:
continue
if isinstance(content, str):
# String content
str_key: tuple[int, int | None] = (choice.index, None)
if str_key in guardrail_map:
if str_key not in already_set:
# First chunk - set the complete guardrailed text
if isinstance(choice, litellm.StreamingChoices):
choice.delta.content = guardrail_map[str_key]
elif isinstance(choice, litellm.Choices):
choice.message.content = guardrail_map[str_key]
already_set[str_key] = True
else:
# Subsequent chunks - clear the content
if isinstance(choice, litellm.StreamingChoices):
choice.delta.content = ""
elif isinstance(choice, litellm.Choices):
choice.message.content = ""
message.content = guardrail_map[str_key] if str_key not in already_set else ""
already_set[str_key] = True
elif isinstance(content, list):
# List content - handle each content item

252
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -11,11 +11,12 @@ import contextlib
import copy
import json
from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence
from itertools import chain
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable
import anyio
from fastapi import HTTPException
from pydantic import TypeAdapter
from pydantic import InstanceOf, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -50,6 +51,7 @@ if TYPE_CHECKING:
A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
GUARDRAIL_NAME: Final = "unified_llm_guardrails"
GUARDRAIL_ADAPTER: Final[TypeAdapter[CustomGuardrail | None]] = TypeAdapter(InstanceOf[CustomGuardrail] | None)
_DECODED_JSON: Final = TypeAdapter(object)
@ -148,6 +150,25 @@ def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
def _choices_owing_text(
buffered: Sequence[object],
emitted_text_per_choice: Mapping[int, str],
held_chars_per_choice: Mapping[int, int],
) -> frozenset[int]:
"""Choices whose guardrailed text has not reached the client yet.
A choice that buffered text but has emitted none of it owes text even when no
processing round has run, so ``held_chars_per_choice`` has no entry for it.
"""
buffered_text: Final = frozenset(
getattr(choice, "index", 0) or 0
for choice in chain.from_iterable(_chunk_choices(chunk) for chunk in buffered)
if getattr(getattr(choice, "delta", None), "content", None)
)
unstarted: Final = frozenset(idx for idx in buffered_text if not emitted_text_per_choice.get(idx))
return unstarted | _held_choices(held_chars_per_choice)
def _recorded_guardrail_information(request_data: _RequestData) -> tuple[StandardLoggingGuardrailInformation, ...]:
_metadata_key, metadata_bucket = get_or_create_metadata_bucket(request_data)
entries: Final = metadata_bucket.get("standard_logging_guardrail_information")
@ -541,12 +562,13 @@ class UnifiedLLMGuardrails(CustomLogger):
self,
*,
reference_chunk: object,
mutated_text_per_choice: dict[int, str],
mutated_text_per_choice: Mapping[int, str],
emitted_text_per_choice: dict[int, str],
holdback_per_choice: dict[int, int],
holdback_per_choice: Mapping[int, int],
finish_reason_per_choice: dict[int, str | None],
held_chars_per_choice: dict[int, int],
is_final: bool,
deliver_finish_reason: bool = True,
) -> ModelResponseStream | None:
"""Build the synthetic chunk carrying the newly-guardrailed deltas.
@ -571,7 +593,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# content+tool_calls chunk whose passthrough suppressed it) still
# needs to reach the client, even if the guardrail returned no text
# to emit. Build a terminator chunk carrying finish_reason per choice.
if is_final and finish_reason_per_choice:
if is_final and deliver_finish_reason and finish_reason_per_choice:
terminator_choices: Final[list[StreamingChoices]] = []
for choice_idx, finish_reason in finish_reason_per_choice.items():
if finish_reason is None:
@ -622,7 +644,7 @@ class UnifiedLLMGuardrails(CustomLogger):
synthetic_choices: Final[list[StreamingChoices]] = []
for choice_idx in mutated_text_per_choice:
delta_text = deltas.get(choice_idx, "")
finish_reason = finish_reason_per_choice.get(choice_idx) if is_final else None
finish_reason = finish_reason_per_choice.get(choice_idx) if is_final and deliver_finish_reason else None
# Skip a choice with nothing to say: no new content and no
# finish_reason to deliver. This avoids emitting an empty delta for an
# already-finished choice (e.g. one that terminated via a passed-through
@ -670,6 +692,7 @@ class UnifiedLLMGuardrails(CustomLogger):
held_chars_per_choice: dict[int, int],
is_final: bool,
terminated: asyncio.Event,
deliver_finish_reason: bool = True,
) -> AsyncGenerator[object, None]:
"""Run one guardrail processing round and emit the resulting diff chunk.
@ -699,6 +722,7 @@ class UnifiedLLMGuardrails(CustomLogger):
finish_reason_per_choice=finish_reason_per_choice,
held_chars_per_choice=held_chars_per_choice,
is_final=is_final,
deliver_finish_reason=deliver_finish_reason,
)
except ModifyResponseException as e:
terminated.set()
@ -760,7 +784,9 @@ class UnifiedLLMGuardrails(CustomLogger):
last_chunk: object | None = None
terminated: Final = asyncio.Event()
def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]:
def _round(
reference_chunk: object, is_final: bool, deliver_finish_reason: bool = True
) -> AsyncGenerator[object, None]:
return self._emit_transform_round(
endpoint_translation=endpoint_translation,
guardrail_to_apply=guardrail_to_apply,
@ -775,50 +801,20 @@ class UnifiedLLMGuardrails(CustomLogger):
held_chars_per_choice=held_chars_per_choice,
is_final=is_final,
terminated=terminated,
deliver_finish_reason=deliver_finish_reason,
)
saw_tool_calls = False
saw_text_content = False
tool_calls_released = False # rebind-ok: set once a raw tool call reaches the client unscanned
end_of_stream_inspection_started = False # rebind-ok: set once the end-of-stream inspection owns the verdict
text_preceded_tool_calls = False # rebind-ok: records text arriving before the first tool call
try:
async for item in response:
# v1 transforms only text. A chunk carrying tool_calls is passed
# through raw so function-calling turns are not dropped, but ONLY
# its tool-call fields are forwarded: content is stripped so any
# response text (in the same delta, or in another choice of an n>1
# chunk) can never bypass the transform. The original chunk is kept
# in responses_so_far so its text is still accumulated + redacted +
# emitted as synthetic deltas, and so the guardrail inspects the
# assembled tool calls at end of stream (see the block inspection
# below), matching block_only. finish_reason rides on the raw
# tool-only chunk, so it is not recorded for the text flush.
if self._chunk_has_tool_calls(item):
text_preceded_tool_calls = text_preceded_tool_calls or (not saw_tool_calls and saw_text_content)
saw_tool_calls = True
responses_so_far.append(item)
last_chunk = item
# Fix #3 — flush accumulated text BEFORE the tool-call
# passthrough. Without this, a stream of text chunks that
# hasn't yet hit a sampled round can be trailed by a
# tool-call chunk carrying finish_reason="tool_calls"; an
# SSE-compliant client stops reading at that finish_reason
# and drops the end-of-stream text flush that would follow.
if saw_text_content:
async for out in _round(item, is_final=False):
yield out
# Fix #1 — pass finish_reason_per_choice into the
# passthrough so a mixed content+tool_call chunk defers its
# finish_reason to the final text terminator (see the
# _tool_call_passthrough_chunk docstring).
tool_only = self._tool_call_passthrough_chunk(
item,
finish_reason_per_choice=finish_reason_per_choice,
held_choices=_held_choices(held_chars_per_choice),
)
responses_yielded.append(tool_only)
tool_calls_released = True
yield tool_only
continue
if self._is_trailing_metadata_chunk(item):
@ -836,39 +832,62 @@ class UnifiedLLMGuardrails(CustomLogger):
# sampled round here would guardrail the same content twice.
if (
not end_of_stream_only
and not saw_tool_calls
and not self._chunk_has_finish_reason(item)
and chunk_counter % sampling_rate == 0
):
async for out in _round(item, is_final=False):
yield out
# v1 does not transform streamed tool calls, but they must still go
# through the guardrail's block decision. Run the block_only inspection
# over the full assembled response so tool calls cannot bypass it.
#
# Pass a deep copy of responses_so_far — the block path routes through
# ``_process_streaming_block_only`` which mutates ``delta.content``
# in-place on the chunk objects it receives. For an n>1 chunk carrying
# text on one choice and tool_calls (with finish_reason) on another,
# ``has_stream_ended`` reads ``choices[0]`` alone and can miss the
# terminal signal, letting the block path rewrite the raw accumulator.
# The subsequent final ``_round`` would then re-read the already-mutated
# text, producing double-application for a non-idempotent guardrail or a
# ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow
# list copy wouldn't help — the mutation is on the chunk objects
# themselves — so we deepcopy.
end_of_stream_inspection_started = True
if saw_tool_calls:
async for out in self._inspect_full_response_for_block(
inspected_responses: Final = copy.deepcopy(responses_so_far)
async for out in self._inspect_full_response(
endpoint_translation=endpoint_translation,
guardrail_to_apply=guardrail_to_apply,
request_data=request_data,
user_api_key_dict=user_api_key_dict,
responses_so_far=copy.deepcopy(responses_so_far),
responses_so_far=inspected_responses,
responses_yielded=responses_yielded,
):
yield out
if text_preceded_tool_calls:
async for out in _round(last_chunk, is_final=True, deliver_finish_reason=False):
yield out
tool_chunks: Final = tuple(
self._tool_call_passthrough_chunk(
buffered_item,
finish_reason_per_choice=finish_reason_per_choice,
held_choices=_choices_owing_text(
inspected_responses, emitted_text_per_choice, held_chars_per_choice
),
)
for buffered_item in inspected_responses
if self._chunk_has_tool_calls(buffered_item)
)
tail_chunks: Final = tuple(
[
chunk
async for chunk in self._checked_stream_tail(
last_chunk=last_chunk,
final_round=_round,
responses_so_far=responses_so_far,
responses_yielded=responses_yielded,
)
]
)
if tail_chunks and isinstance(tail_chunks[-1], _StreamTerminated):
for error_chunk in tail_chunks[:-1]:
yield error_chunk
return
for tool_only in tool_chunks:
responses_yielded.append(tool_only)
yield tool_only
for tail_chunk in tail_chunks:
yield tail_chunk
return
async for out in self._emit_stream_tail(
last_chunk=last_chunk,
final_round=_round,
@ -878,37 +897,25 @@ class UnifiedLLMGuardrails(CustomLogger):
yield out
except _StreamTerminated:
return
except (GeneratorExit, asyncio.CancelledError):
await self._scan_uninspected_tool_calls_after_disconnect(
uninspected=tool_calls_released and not end_of_stream_inspection_started and not terminated.is_set(),
endpoint_translation=endpoint_translation,
responses_released=responses_yielded,
guardrail_to_apply=guardrail_to_apply,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
raise
@staticmethod
async def _scan_uninspected_tool_calls_after_disconnect(
async def _checked_stream_tail(
self,
*,
uninspected: bool,
endpoint_translation: _EndpointTranslation,
responses_released: Sequence[object],
guardrail_to_apply: CustomGuardrail,
user_api_key_dict: UserAPIKeyAuth,
request_data: _RequestData,
) -> None:
if not uninspected:
return
await UnifiedLLMGuardrails._scan_released_stream_after_disconnect(
endpoint_translation=endpoint_translation,
responses_released=responses_released,
last_scan_key=None,
guardrail_to_apply=guardrail_to_apply,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
last_chunk: object | None,
final_round: Callable[[object, bool], AsyncGenerator[object, None]],
responses_so_far: Sequence[object],
responses_yielded: list[object],
) -> AsyncGenerator[object, None]:
try:
async for chunk in self._emit_stream_tail(
last_chunk=last_chunk,
final_round=final_round,
responses_so_far=responses_so_far,
responses_yielded=responses_yielded,
):
yield chunk
except _StreamTerminated as exc:
yield exc
async def _emit_stream_tail(
self,
@ -927,29 +934,21 @@ class UnifiedLLMGuardrails(CustomLogger):
responses_yielded.append(trailing)
yield trailing
async def _inspect_full_response_for_block(
async def _inspect_full_response(
self,
*,
endpoint_translation: _EndpointTranslation,
guardrail_to_apply: CustomGuardrail,
request_data: dict,
request_data: _RequestData,
user_api_key_dict: UserAPIKeyAuth,
responses_so_far: Sequence[object],
responses_yielded: Sequence[object],
) -> AsyncGenerator[object, None]:
"""Run the block-only guardrail inspection over the full assembled
response (text + tool calls) so nothing bypasses the block decision.
The guardrail's returned transforms are discarded here (v1 does not
transform tool calls); only its block decision matters. A block is
surfaced the same way as elsewhere: ModifyResponseException terminates the
stream via the shared block handler; a GenericGuardrailAPI block raises and
propagates, matching block_only.
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
try:
with anyio.CancelScope(shield=bool(responses_yielded)):
with anyio.CancelScope(shield=True):
await endpoint_translation.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail_to_apply,
@ -957,7 +956,10 @@ class UnifiedLLMGuardrails(CustomLogger):
user_api_key_dict=user_api_key_dict,
request_data=request_data,
stream_transform_sink=None,
deliver_ended_stream_rewrites=True,
)
except UndeliverableStreamRewrite as exc:
raise HTTPException(status_code=400, detail="Guardrail stream rewrite could not be applied") from exc
except ModifyResponseException as e:
if e.original_response is None:
e.original_response = responses_so_far
@ -1050,6 +1052,10 @@ class UnifiedLLMGuardrails(CustomLogger):
finish_reason=passthrough_finish,
)
)
if isinstance(item, ModelResponseStream):
passthrough: Final = copy.deepcopy(item)
passthrough.choices = synthetic_choices
return passthrough
return ModelResponseStream(
id=getattr(item, "id", None),
created=getattr(item, "created", None),
@ -1141,12 +1147,15 @@ class UnifiedLLMGuardrails(CustomLogger):
# litellm.integrations.custom_guardrail.
from litellm.integrations.custom_guardrail import ModifyResponseException
if guardrail_to_apply is None:
guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
resolved_guardrail: Final[CustomGuardrail | None] = (
guardrail_to_apply
if guardrail_to_apply is not None
else GUARDRAIL_ADAPTER.validate_python(request_data.pop("guardrail_to_apply", None))
)
typed_request_data: Final[_RequestData] = request_data
def _streaming_setting(name: str, default: object) -> object:
return self.resolve_streaming_flag(guardrail_to_apply, name, default)
return self.resolve_streaming_flag(resolved_guardrail, name, default)
def _streaming_flag(name: str, default: object) -> Any:
return _streaming_setting(name, default)
@ -1172,14 +1181,14 @@ class UnifiedLLMGuardrails(CustomLogger):
if (
buffer_until_moderated
and guardrail_to_apply is not None
and getattr(guardrail_to_apply, "mask_response_content", False)
and resolved_guardrail is not None
and getattr(resolved_guardrail, "mask_response_content", False)
):
verbose_proxy_logger.warning(
"UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s "
"because mask_response_content=True -- buffered replay would release "
"unredacted original chunks instead of the moderated output.",
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
buffer_until_moderated = False
@ -1193,16 +1202,16 @@ class UnifiedLLMGuardrails(CustomLogger):
buffer_until_moderated or end_of_stream_only
)
if guardrail_to_apply is None:
if resolved_guardrail is None:
async for item in response:
yield item
return
event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call
if guardrail_to_apply.should_run_guardrail(data=request_data, event_type=event_type) is not True:
if resolved_guardrail.should_run_guardrail(data=request_data, event_type=event_type) is not True:
verbose_proxy_logger.debug(
"UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s",
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
async for item in response:
yield item
@ -1222,7 +1231,7 @@ class UnifiedLLMGuardrails(CustomLogger):
if transform_call_type is not None:
async with contextlib.aclosing(
self._run_incremental_transform_stream(
guardrail_to_apply=guardrail_to_apply,
guardrail_to_apply=resolved_guardrail,
response=response,
request_data=typed_request_data,
user_api_key_dict=user_api_key_dict,
@ -1239,7 +1248,7 @@ class UnifiedLLMGuardrails(CustomLogger):
"UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported "
"for the OpenAI chat completions streaming path with a resolvable request route; "
"falling back to block_only for %s",
getattr(guardrail_to_apply, "guardrail_name", None),
getattr(resolved_guardrail, "guardrail_name", None),
)
# Infer call type from first chunk
@ -1333,7 +1342,7 @@ class UnifiedLLMGuardrails(CustomLogger):
verbose_proxy_logger.debug(
"Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
chunk_counter,
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
if buffer_until_moderated:
if hold_window:
@ -1353,14 +1362,14 @@ class UnifiedLLMGuardrails(CustomLogger):
"Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
chunk_counter,
sampling_rate,
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
if scan_in_background:
scanned_responses.extend(copy.deepcopy(responses_so_far[len(scanned_responses) :]))
scan = endpoint_translation.process_output_streaming_response(
responses_so_far=list(scanned_responses) if scan_in_background else responses_so_far,
guardrail_to_apply=guardrail_to_apply,
guardrail_to_apply=resolved_guardrail,
litellm_logging_obj=request_data.get("litellm_logging_obj"),
user_api_key_dict=user_api_key_dict,
request_data=request_data,
@ -1421,7 +1430,7 @@ class UnifiedLLMGuardrails(CustomLogger):
verbose_proxy_logger.debug(
"Holding %s buffered chunks for guardrail %s: this round could not scan the whole window",
len(withheld_items),
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
withheld_items[:] = original_items
continue
@ -1441,7 +1450,7 @@ class UnifiedLLMGuardrails(CustomLogger):
verbose_proxy_logger.debug(
"Processing final streaming response with all %s chunks for guardrail %s",
len(responses_so_far),
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
endpoint_translation = mappings[CallTypes(call_type)]()
@ -1461,12 +1470,12 @@ class UnifiedLLMGuardrails(CustomLogger):
verdict_settled = True
if (
_is_redundant_scan(end_scan_key, last_scan_key)
and not _owes_ended_scan(endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key)
and not _owes_ended_scan(endpoint_translation, resolved_guardrail, responses_so_far, last_scan_key)
and not background_scan_failed
):
verbose_proxy_logger.debug(
"Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
guardrail_to_apply.guardrail_name,
resolved_guardrail.guardrail_name,
)
for buffered_item in buffered_items or ():
yield buffered_item
@ -1481,7 +1490,7 @@ class UnifiedLLMGuardrails(CustomLogger):
pending_scan.result() # a failed background scan raises into the handlers below
await endpoint_translation.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail_to_apply,
guardrail_to_apply=resolved_guardrail,
litellm_logging_obj=request_data.get("litellm_logging_obj"),
user_api_key_dict=user_api_key_dict,
request_data=request_data,
@ -1526,17 +1535,12 @@ class UnifiedLLMGuardrails(CustomLogger):
if isinstance(pending_scan.exception(), (ModifyResponseException, HTTPException)):
verdict_settled = True # rebind-ok: the background scan already reached a verdict
translation_class: Final = None if call_type is None else mappings.get(CallTypes(call_type))
if (
chunks_yielded
and not verdict_settled
and translation_class is not None
and isinstance(guardrail_to_apply, CustomGuardrail)
):
if chunks_yielded and not verdict_settled and translation_class is not None:
await self._scan_released_stream_after_disconnect(
endpoint_translation=translation_class(),
responses_released=responses_yielded,
last_scan_key=last_scan_key,
guardrail_to_apply=guardrail_to_apply,
guardrail_to_apply=resolved_guardrail,
user_api_key_dict=user_api_key_dict,
request_data=typed_request_data,
)

4
litellm/types/utils.py
View file

@ -3645,6 +3645,7 @@ AUDIT_GUARDRAIL_FIELDS: Final[frozenset[str]] = frozenset(
"end_time",
"duration",
"masked_entity_count",
"choice_index",
"guardrail_id",
"policy_template",
"detection_method",
@ -3685,6 +3686,9 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False):
}
"""
choice_index: ReadOnly[int | None]
"""Which choice of an n > 1 response this entry inspected, when the guardrail ran per choice."""
guardrail_id: str | None
"""Unique identifier for the guardrail configuration, e.g. 'gd-eu-pii-001'"""

13
tests/unit/integrations/otel/test_otel_v2_emitter.py
View file

@ -344,6 +344,19 @@ def test_guardrail_block_span_is_error_and_carries_verdict():
assert a[LiteLLM.GUARDRAIL_MASKED_ENTITY_COUNT] == 2
def test_guardrail_span_carries_the_choice_index():
"""A per-choice guardrail entry must say on the span which choice it inspected."""
engine, exporter = _engine()
engine.emit(
SpanRole.GUARDRAIL,
GuardrailSpanData.from_logging_entry(
{"guardrail_name": "g", "guardrail_status": "success", "choice_index": 1}
),
)
(span,) = exporter.get_finished_spans()
assert span.attributes[LiteLLM.GUARDRAIL_CHOICE_INDEX] == 1
def test_guardrail_success_span_is_unset():
"""On success the status is left UNSET (semconv default) — not forced OK."""
engine, exporter = _engine()

543
tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -7,7 +7,7 @@ with guardrail transformations, including tool calls.
import json
from collections.abc import Mapping
from typing import Any, Literal, Optional
from typing import Any, Final, Literal, Optional
import pytest
@ -1146,10 +1146,12 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
)
assert result is chunks
fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]]
assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""]
assert fragments[0][0].function.name == "lookup_fruit"
assert fragments[0][0].id == "call_1"
fragments = [call for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()]
arguments = [fragment.function.arguments for fragment in fragments]
assert "".join(arguments) == '{"fruit": "PERSIMMON"}'
assert arguments.count("") <= 1
assert fragments[0].function.name == "lookup_fruit"
assert fragments[0].id == "call_1"
assert chunks[3].choices[0].delta.tool_calls is None
assert chunks[3].choices[0].finish_reason == "tool_calls"
@ -1171,8 +1173,10 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
deliver_ended_stream_rewrites=True,
)
fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]]
assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None]
fragments = [call for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()]
assert [fragment.function.name for fragment in fragments if fragment.function.name] == [
"lookup_fruit_reviewed"
]
assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"}
assert fragments[0].id == "call_1"
@ -1395,29 +1399,524 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
)
assert [
(tool_call["id"], tool_call["function"]["arguments"])
for tool_call in guardrail.seen_inputs[-1]["tool_calls"]
] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')]
[(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]]
for inputs in guardrail.seen_inputs
] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]]
@pytest.mark.asyncio
async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self):
@pytest.mark.parametrize("transform", [False, True])
async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None:
from fastapi import HTTPException
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class ContextGuardrail(CustomGuardrail):
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
response: Final = request_data["response"]
assert isinstance(response, ModelResponse)
request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index)
if response.choices[0].message.content == "forbidden":
raise HTTPException(status_code=400, detail="later choice blocked")
return inputs
chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
index=index, delta=Delta(content=text, tool_calls=[{
"index": 0, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": "{}"},
}]), finish_reason="tool_calls",
) for index, text in enumerate(("allowed", "forbidden"))])]
request_data: Final = {"metadata": {"trace": "retained"}}
with pytest.raises(HTTPException, match="later choice blocked"):
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"),
request_data=request_data, deliver_ended_stream_rewrites=True,
stream_transform_sink=StreamTransformSink() if transform else None,
)
assert request_data["visited_choices"] == (0, 1)
assert request_data["metadata"]["trace"] == "retained"
assert "response" not in request_data
assert "responses" not in request_data
@pytest.mark.asyncio
async def test_undeliverable_rewrite_still_inspects_the_remaining_choices(self) -> None:
"""An undeliverable rewrite on one choice must not leave later choices uninspected.
The caller releases the original chunks when it catches UndeliverableStreamRewrite, so
bailing out of the loop would let a request with n > 1 pair a redactable choice with a
prohibited one and have the whole response delivered unchecked.
"""
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_tool_call_stream_chunks()
class UndeliverableOnFirstChoice(CustomGuardrail):
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
response: Final = request_data["response"]
assert isinstance(response, ModelResponse)
index: Final = response.choices[0].index
request_data["visited_choices"] = (*request_data.get("visited_choices", ()), index)
if index == 0:
raise UndeliverableStreamRewrite("undeliverable", "tool_call_mismatch")
return inputs
with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised:
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
index=index, delta=Delta(content=text, tool_calls=[{
"index": 0, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": "{}"},
}]), finish_reason="tool_calls",
) for index, text in enumerate(("redacted", "prohibited"))])]
request_data: Final = {}
with pytest.raises(UndeliverableStreamRewrite):
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=UndeliverableOnFirstChoice(guardrail_name="undeliverable"),
request_data=request_data, deliver_ended_stream_rewrites=True,
)
assert request_data["visited_choices"] == (0, 1)
@pytest.mark.asyncio
async def test_undeliverable_rewrite_in_transform_mode_still_inspects_remaining_choices(self) -> None:
"""The transform path runs its own per-choice loop and needs the same rule."""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class UndeliverableOnFirstChoice(CustomGuardrail):
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
response: Final = request_data["response"]
assert isinstance(response, ModelResponse)
index: Final = response.choices[0].index
request_data["visited_choices"] = (*request_data.get("visited_choices", ()), index)
if index == 0:
raise UndeliverableStreamRewrite("undeliverable", "tool_call_mismatch")
return inputs
chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
index=index, delta=Delta(content=text, tool_calls=[{
"index": 0, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": "{}"},
}]), finish_reason="tool_calls",
) for index, text in enumerate(("redacted", "prohibited"))])]
request_data: Final = {}
with pytest.raises(UndeliverableStreamRewrite):
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=UndeliverableOnFirstChoice(guardrail_name="undeliverable"),
request_data=request_data, deliver_ended_stream_rewrites=True,
stream_transform_sink=StreamTransformSink(),
)
assert request_data["visited_choices"] == (0, 1)
@pytest.mark.asyncio
async def test_buffered_stream_without_finish_reason_takes_the_ended_path(self) -> None:
"""A buffered stream has no finish_reason yet, so the ended path must key off the opt-in.
Dropping ``deliver_ended_stream_rewrites`` from that condition leaves the buffered
caller on the in-flight path, where rewrites are never written back.
"""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class PathProbe(CustomGuardrail):
saw_assembled_response: bool = False
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).saw_assembled_response = isinstance(request_data.get("response"), ModelResponse)
return inputs
chunks: Final = [
ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content="secret"), finish_reason=None)]
)
]
assert OpenAIChatCompletionsHandler()._first_choice_has_finished(chunks) is False
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=PathProbe(guardrail_name="probe"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert PathProbe.saw_assembled_response is True
@pytest.mark.asyncio
async def test_diff_mode_keeps_each_choice_holdback_with_two_choices(self) -> None:
"""The per-choice sinks are merged, so an n=2 diff round must keep both holdbacks.
Reporting an empty holdback here would release text the guardrail asked to withhold.
"""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class HoldingGuardrail(CustomGuardrail):
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
texts: Final = inputs.get("texts") or ()
return {**inputs, "stream_holdback_chars": [len(text) for text in texts]}
chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
index=index, delta=Delta(content=text, tool_calls=[{
"index": 0, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": "{}"},
}]), finish_reason="tool_calls",
) for index, text in enumerate(("alpha", "bravo-long"))])]
sink: Final = StreamTransformSink()
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=HoldingGuardrail(guardrail_name="holding"),
litellm_logging_obj=None,
stream_transform_sink=sink,
)
assert sink.holdback_per_choice == {0: len("alpha"), 1: len("bravo-long")}
@pytest.mark.asyncio
async def test_diff_mode_withholds_tool_calls_until_their_arguments_are_complete(self) -> None:
"""Mid-round the assembled arguments are a truncated JSON fragment, so they are not sent.
Sending them would hand the guardrail something like `{"ci`, which main never did.
"""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class RecordingGuardrail(CustomGuardrail):
seen: list[object] = []
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).seen.append(inputs.get("tool_calls"))
return inputs
def chunk(text: str, arguments: str, finish_reason: str | None) -> ModelResponseStream:
return ModelResponseStream(choices=[StreamingChoices(
index=0, delta=Delta(content=text, tool_calls=[{
"index": 0, "id": "call_0", "type": "function",
"function": {"name": "weather", "arguments": arguments},
}]), finish_reason=finish_reason,
)])
RecordingGuardrail.seen = []
handler: Final = OpenAIChatCompletionsHandler()
guardrail: Final = RecordingGuardrail(guardrail_name="recording")
partial: Final = [chunk("looking", '{"ci', None)]
await handler.process_output_streaming_response(
responses_so_far=partial, guardrail_to_apply=guardrail,
litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(),
)
assert RecordingGuardrail.seen == [None]
RecordingGuardrail.seen = []
complete: Final = [chunk("looking", '{"ci', None), chunk(" it up", 'ty": "Paris"}', "tool_calls")]
await handler.process_output_streaming_response(
responses_so_far=complete, guardrail_to_apply=guardrail,
litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(),
)
sent: Final = RecordingGuardrail.seen[-1]
assert sent is not None
assert [call["function"]["arguments"] for call in sent] == ['{"city": "Paris"}']
@pytest.mark.asyncio
async def test_diff_mode_withholds_parallel_tool_calls_until_every_one_is_complete(self) -> None:
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class RecordingGuardrail(CustomGuardrail):
seen: list[object] = []
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).seen.append(inputs.get("tool_calls"))
return inputs
def call(index: int, arguments: str) -> dict[str, object]:
return {
"index": index, "id": f"call_{index}", "type": "function",
"function": {"name": "weather", "arguments": arguments},
}
RecordingGuardrail.seen = []
handler: Final = OpenAIChatCompletionsHandler()
guardrail: Final = RecordingGuardrail(guardrail_name="recording")
one_still_empty: Final = [ModelResponseStream(choices=[StreamingChoices(
index=0, delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, "")]),
finish_reason=None,
)])]
await handler.process_output_streaming_response(
responses_so_far=one_still_empty, guardrail_to_apply=guardrail,
litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(),
)
assert RecordingGuardrail.seen == [None]
RecordingGuardrail.seen = []
both_complete: Final = [ModelResponseStream(choices=[StreamingChoices(
index=0,
delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, '{"city": "Rome"}')]),
finish_reason="tool_calls",
)])]
await handler.process_output_streaming_response(
responses_so_far=both_complete, guardrail_to_apply=guardrail,
litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(),
)
sent: Final = RecordingGuardrail.seen[-1]
assert sent is not None
assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}']
@pytest.mark.asyncio
async def test_a_buffered_stream_that_carried_no_choices_is_not_an_error(self) -> None:
"""A stream can end carrying only usage, and a policy step still runs over it."""
from litellm.types.utils import ModelResponseStream
chunks: Final = [ModelResponseStream(choices=[], usage={"prompt_tokens": 3, "completion_tokens": 0,
"total_tokens": 3})]
result: Final = await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
request_data={},
deliver_ended_stream_rewrites=True,
)
assert result is chunks
@pytest.mark.asyncio
async def test_a_round_with_nothing_to_judge_costs_one_call_at_n2(self) -> None:
"""Per-choice rounds buy tool-call context. While no call is complete there is none to buy,
and charging a vendor once per choice for every sampled round buys nothing."""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class CountingGuardrail(CustomGuardrail):
calls: list[object] = []
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).calls.append(inputs.get("texts"))
return inputs
def choice(index: int, arguments: str) -> StreamingChoices:
return StreamingChoices(
index=index,
delta=Delta(content="looking", tool_calls=[
{"index": 0, "id": f"call_{index}", "type": "function",
"function": {"name": "weather", "arguments": arguments}},
]),
finish_reason=None,
)
assert raised.value.guardrail_name == "test"
assert raised.value.reason == (
"the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams"
CountingGuardrail.calls = []
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=[ModelResponseStream(choices=[choice(0, '{"city": "Pa'), choice(1, '{"ci')])],
guardrail_to_apply=CountingGuardrail(guardrail_name="counting"),
litellm_logging_obj=None, request_data={},
stream_transform_sink=StreamTransformSink(),
)
assert len(CountingGuardrail.calls) == 1
def test_an_id_in_flight_on_one_choice_does_not_prune_another(self) -> None:
"""Providers number tool calls per choice, so the same id on another choice is another call,
and one still in flight must not take a finished call down with it."""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
def choice(index: int, arguments: str) -> StreamingChoices:
return StreamingChoices(
index=index,
delta=Delta(content="looking", tool_calls=[
{"index": 0, "id": "call_shared", "type": "function",
"function": {"name": "weather", "arguments": arguments}},
]),
finish_reason="tool_calls",
)
handler: Final = OpenAIChatCompletionsHandler()
chunks: Final = [ModelResponseStream(choices=[choice(0, '{"city": "Pa'), choice(1, '{"city": "Rome"}')])]
assembled: Final = handler._rebuild_ended_stream_per_choice(chunks, None)
pruned: Final = handler._without_unfinished_tool_calls(assembled, chunks)
assert [
[call.function.arguments for call in single.message.tool_calls or ()] for single in pruned.choices
] == [[], ['{"city": "Rome"}']]
@pytest.mark.asyncio
async def test_partial_arguments_stay_out_of_what_each_choice_inspects(self) -> None:
"""Half a JSON object is not a call a guardrail can judge, at any number of choices."""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class RecordingGuardrail(CustomGuardrail):
seen: list[object] = []
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).seen.append(request_data.get("response"))
return inputs
def choice(index: int, arguments: str, label: str) -> StreamingChoices:
return StreamingChoices(
index=index,
delta=Delta(content="looking", tool_calls=[
{"index": 0, "id": "call_" + label, "type": "function",
"function": {"name": "weather", "arguments": arguments}},
]),
finish_reason="tool_calls",
)
RecordingGuardrail.seen = []
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=[ModelResponseStream(choices=[
choice(0, '{"city": "Pa', "partial"), choice(1, '{"city": "Rome"}', "rome"),
])],
guardrail_to_apply=RecordingGuardrail(guardrail_name="recording"),
litellm_logging_obj=None, request_data={},
stream_transform_sink=StreamTransformSink(),
)
arguments: Final = tuple(
tool_call.function.arguments
for response in RecordingGuardrail.seen
for single in getattr(response, "choices", None) or ()
for tool_call in single.message.tool_calls or ()
)
assert arguments == ('{"city": "Rome"}',)
@pytest.mark.asyncio
async def test_a_caller_cannot_set_the_choice_index_label(self) -> None:
"""The label says which choice the round inspected, so only the round may set it."""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class LabelReadingGuardrail(CustomGuardrail):
seen: list[object] = []
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).seen.append(request_data.get("guardrail_choice_index"))
return inputs
LabelReadingGuardrail.seen = []
request_data: Final[dict[str, object]] = {"model": "gpt-4", "guardrail_choice_index": 7}
chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
index=0,
delta=Delta(content="looking", tool_calls=[
{"index": 0, "id": "call_0", "type": "function",
"function": {"name": "weather", "arguments": '{"city": "Paris"}'}},
]),
finish_reason="tool_calls",
)])]
await OpenAIChatCompletionsHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=LabelReadingGuardrail(guardrail_name="label"),
litellm_logging_obj=None,
request_data=request_data,
stream_transform_sink=StreamTransformSink(),
)
assert LabelReadingGuardrail.seen == [0]
assert request_data["guardrail_choice_index"] == 7
@pytest.mark.asyncio
async def test_a_transform_round_leaves_the_callers_request_data_as_it_found_it(self) -> None:
"""The round owns these keys only while it runs; what it leaves behind is read by later hooks."""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler: Final = OpenAIChatCompletionsHandler()
request_data: Final[dict[str, object]] = {"model": "gpt-4"}
chunks: Final = [ModelResponseStream(choices=[StreamingChoices(
index=0,
delta=Delta(content="looking", tool_calls=[
{"index": 0, "id": "call_0", "type": "function",
"function": {"name": "weather", "arguments": '{"city": "Paris"}'}},
]),
finish_reason="tool_calls",
)])]
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
request_data=request_data,
stream_transform_sink=StreamTransformSink(),
)
assert "response" not in request_data
assert "responses" not in request_data
assert "guardrail_choice_index" not in request_data
@pytest.mark.asyncio
async def test_an_unfinished_parallel_call_is_kept_out_of_what_the_guardrail_reads(self) -> None:
"""A call that only announced its name must not reach the guardrail, by either route."""
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
class RecordingGuardrail(CustomGuardrail):
seen_calls: list[object] = []
seen_response: list[object] = []
async def apply_guardrail(
self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object],
input_type: Literal["request", "response"], logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
type(self).seen_calls.append(inputs.get("tool_calls"))
type(self).seen_response.append(request_data.get("response"))
return inputs
def call(index: int, arguments: str) -> dict[str, object]:
return {
"index": index, "id": f"call_{index}", "type": "function",
"function": {"name": "weather", "arguments": arguments},
}
RecordingGuardrail.seen_calls = []
RecordingGuardrail.seen_response = []
handler: Final = OpenAIChatCompletionsHandler()
request_data: Final[dict[str, object]] = {}
ended_with_one_unfinished: Final = [ModelResponseStream(choices=[StreamingChoices(
index=0, delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, "")]),
finish_reason="tool_calls",
)])]
await handler.process_output_streaming_response(
responses_so_far=ended_with_one_unfinished,
guardrail_to_apply=RecordingGuardrail(guardrail_name="recording"),
litellm_logging_obj=None, request_data=request_data,
stream_transform_sink=StreamTransformSink(),
)
assert RecordingGuardrail.seen_calls == [None]
inspected: Final = RecordingGuardrail.seen_response[-1]
arguments: Final = tuple(
tool_call.function.arguments
for choice in getattr(inspected, "choices", None) or ()
for tool_call in choice.message.tool_calls or ()
)
assert arguments == ('{"city": "Paris"}',)
@pytest.mark.asyncio
async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None:
handler: Final = OpenAIChatCompletionsHandler()
chunks: Final = self._two_choice_tool_call_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
arguments: Final = tuple(
"".join(call.function.arguments for chunk in chunks for choice in chunk.choices
if choice.index == index for call in choice.delta.tool_calls or ())
for index in range(2)
)
assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}')
@pytest.mark.asyncio
async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):

353
tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -10,6 +10,8 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal
from pydantic import TypeAdapter
import anyio
import pytest
@ -53,6 +55,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
ChatCompletionMessageToolCall,
Delta,
GenericGuardrailAPIInputs,
ModelResponse,
@ -1011,6 +1014,42 @@ def _delta_text(item):
return item.choices[0].delta.content or ""
class _ToolRedactingGuardrail(CustomGuardrail):
def __init__(self, require_tool_context: bool = False) -> None:
super().__init__(guardrail_name="tool-redactor", event_hook=GuardrailEventHooks.post_call, default_on=True)
self.streaming_transform_mode = "incremental_diff"
self.streaming_sampling_rate = 1
self.require_tool_context = require_tool_context
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: object | None = None,
) -> GenericGuardrailAPIInputs:
calls: Final = TypeAdapter(tuple[ChatCompletionMessageToolCall, ...]).validate_python(
inputs.get("tool_calls", ())
)
input_texts: Final = inputs.get("texts", ())
texts: Final = tuple(
"checked:" + text.replace("SECRET", "MASKED")
if not self.require_tool_context or (calls and index == len(input_texts) - 1) else text
for index, text in enumerate(input_texts)
)
return {
**inputs,
"texts": list(texts),
"stream_holdback_chars": [len(text) for text in texts],
"tool_calls": [
call.model_copy(update={"function": call.function.model_copy(update={
"arguments": call.function.arguments.replace("SECRET", "MASKED"),
})})
for call in calls
],
}
class TestStreamingTransform:
"""Streaming text-transformation (incremental_diff) path on the OpenAI chat
completions streaming surface."""
@ -1019,6 +1058,103 @@ class TestStreamingTransform:
def _use_openai_handler_mapping(self, monkeypatch):
_patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler})
@pytest.mark.asyncio
@pytest.mark.parametrize("require_tool_context", [False, True])
@pytest.mark.parametrize("include_text", [False, True])
@pytest.mark.parametrize("tool_count", [1, 2])
async def test_buffered_tool_arguments_are_rewritten_before_delivery(
self, include_text: bool, tool_count: int, require_tool_context: bool
) -> None:
chunks: Final = (
*([_stream_chunk("hello SECRET")] if include_text else []),
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
{"index": index, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"SEC'}}
for index in range(tool_count)
]))]),
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
{"index": index, "function": {"arguments": 'RET"}'}} for index in range(tool_count)
]), finish_reason="tool_calls")]),
ModelResponseStream(choices=[], usage={"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}),
)
out: Final = await _drive_stream(
UnifiedLLMGuardrails(), _ToolRedactingGuardrail(require_tool_context=require_tool_context), chunks
)
calls: Final = tuple(
call for chunk in out for choice in chunk.choices for call in choice.delta.tool_calls or ()
)
for index in range(tool_count):
arguments: Final = "".join(call.function.arguments or "" for call in calls if call.index == index)
assert arguments == '{"contact":"MASKED"}'
assert next(call.id for call in calls if call.index == index and call.id) == f"call_{index}"
assert "".join(_delta_text(chunk) for chunk in out) == ("checked:hello MASKED" if include_text else "")
assert any(choice.finish_reason == "tool_calls" for chunk in out for choice in chunk.choices)
assert out[-1].usage.total_tokens == 18
assert all("SECRET" not in chunk.model_dump_json() for chunk in out)
@pytest.mark.asyncio
@pytest.mark.parametrize("tool_choice_index", [0, 1])
async def test_tool_dependent_text_rewrites_keep_choice_context(self, tool_choice_index: int) -> None:
chunks: Final = (
ModelResponseStream(choices=[StreamingChoices(
index=index, delta=Delta(content="SECRET" if index == tool_choice_index else "plain"),
) for index in (1, 0)]),
ModelResponseStream(choices=[StreamingChoices(
index=tool_choice_index,
delta=Delta(tool_calls=[{
"index": 0, "id": "call_context", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"SECRET"}'},
}]), finish_reason="tool_calls",
)]),
)
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(True), chunks)
for index in (0, 1):
text: Final = "".join(
choice.delta.content or "" for chunk in out for choice in chunk.choices if choice.index == index
)
assert text == ("checked:MASKED" if index == tool_choice_index else "plain")
assert all("SECRET" not in chunk.model_dump_json() for chunk in out)
@pytest.mark.asyncio
async def test_tool_rewrites_keep_completion_choices_separate(self) -> None:
async def response() -> AsyncIterator[ModelResponseStream]:
yield ModelResponseStream(choices=[StreamingChoices(
index=index,
delta=Delta(tool_calls=[{"index": 0, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"SECRET"}'}}]),
finish_reason="tool_calls",
) for index in range(2)])
iterator: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
response=response(), request_data={"guardrail_to_apply": _ToolRedactingGuardrail()},
)
chunks: Final = tuple([chunk async for chunk in iterator])
calls: Final = tuple(
(choice.index, call.id, call.function.arguments)
for chunk in chunks for choice in chunk.choices for call in choice.delta.tool_calls or ()
)
assert calls == ((0, "call_0", '{"contact":"MASKED"}'), (1, "call_1", '{"contact":"MASKED"}'))
@pytest.mark.asyncio
async def test_undeliverable_rewrite_is_a_closed_failure(self) -> None:
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
class UndeliverableTranslation(OpenAIChatCompletionsHandler):
async def process_output_streaming_response(
self, *args: object, **kwargs: object
) -> list[ModelResponseStream]:
raise UndeliverableStreamRewrite("tool-redactor", "unmappable tool fragments")
iterator: Final = UnifiedLLMGuardrails()._inspect_full_response(
endpoint_translation=UndeliverableTranslation(), guardrail_to_apply=_ToolRedactingGuardrail(),
request_data={}, user_api_key_dict=UserAPIKeyAuth(), responses_so_far=(), responses_yielded=(),
)
with pytest.raises(unified_module.HTTPException) as error:
await anext(iterator)
assert error.value.status_code == 400
assert error.value.detail == "Guardrail stream rewrite could not be applied"
@pytest.mark.asyncio
async def test_block_only_drops_text_rewrites(self):
"""Default block_only: the guardrail's uppercasing never reaches the
@ -1038,6 +1174,179 @@ class TestStreamingTransform:
assert streamed == "hello world"
assert streamed != streamed.upper()
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_rewritten_tool_arguments_do_not_leave_empty_fragments_on_the_wire(self):
"""Blanking the trailing fragments pads the stream with meaningless tool-call deltas."""
def fragment(
arguments: str, announce: bool = False, finish_reason: str | None = None
) -> ModelResponseStream:
call: Final[dict[str, object]] = {"index": 0, "function": {"arguments": arguments}}
announced: Final = (
{**call, "id": "call_0", "type": "function", "function": {"name": "lookup", "arguments": arguments}}
if announce
else call
)
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(content=None, role="assistant", tool_calls=[announced]),
finish_reason=finish_reason,
)
],
)
chunks: Final = [
_stream_chunk("looking"),
fragment("", announce=True),
fragment('{"q": "SE'),
fragment('CRET"}', finish_reason="tool_calls"),
]
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks)
delivered_arguments: Final = tuple(
call.function.arguments
for item in out
for choice in getattr(item, "choices", ()) or ()
for call in getattr(choice.delta, "tool_calls", None) or ()
)
assert "".join(delivered_arguments) == '{"q": "MASKED"}'
assert delivered_arguments.count("") <= 1
@pytest.mark.asyncio
async def test_a_tool_call_streamed_before_text_is_delivered_before_that_text(self) -> None:
"""The replay must keep the order the model chose between tool calls and text."""
chunks: Final = (
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
{"index": 0, "id": "call_0", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"ok"}'}},
]))]),
_stream_chunk("hello SECRET", finish_reason="stop"),
)
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks)
kinds: Final = tuple(
"tool" if choice.delta.tool_calls else "text"
for chunk in out
for choice in chunk.choices
if choice.delta.tool_calls or choice.delta.content
)
assert kinds.index("tool") < kinds.index("text")
@pytest.mark.asyncio
async def test_text_streamed_before_a_tool_call_is_delivered_before_it(self) -> None:
"""The replay must not move a tool call ahead of text the model sent first."""
chunks: Final = (
_stream_chunk("hello "),
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
{"index": 0, "id": "call_0", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"ok"}'}},
]))]),
_stream_chunk("SECRET", finish_reason="stop"),
)
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks)
kinds: Final = tuple(
"tool" if choice.delta.tool_calls else "text"
for chunk in out
for choice in chunk.choices
if choice.delta.tool_calls or choice.delta.content
)
assert kinds[0] == "text"
assert "tool" in kinds
assert "".join(_delta_text(chunk) for chunk in out) == "checked:hello MASKED"
@pytest.mark.asyncio
async def test_tool_call_replay_keeps_the_provider_fields_of_its_chunk(self) -> None:
"""A replayed tool-call chunk must not lose the provider metadata the model sent with it."""
chunks: Final = (
_stream_chunk("hello SECRET"),
ModelResponseStream(
id="chatcmpl-provider",
created=1234567890,
model="gpt-4",
system_fingerprint="fp_provider",
provider_specific_fields={"service_tier": "scale"},
choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
{"index": 0, "id": "call_0", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"ok"}'}},
]), finish_reason="tool_calls")],
),
)
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks)
replayed: Final = tuple(
chunk for chunk in out for choice in chunk.choices if choice.delta.tool_calls
)
assert replayed
assert all(chunk.system_fingerprint == "fp_provider" for chunk in replayed)
assert all(chunk.provider_specific_fields == {"service_tier": "scale"} for chunk in replayed)
@pytest.mark.asyncio
async def test_finish_on_a_tool_chunk_never_precedes_the_rewritten_text(self) -> None:
"""A client that stops at finish_reason must already have the guardrailed text."""
def tool(index: int, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[
{"index": index, "id": f"call_{index}", "type": "function",
"function": {"name": "contact", "arguments": '{"contact":"ok"}'}},
]), finish_reason=finish_reason)])
chunks: Final = (tool(0), _stream_chunk("hello SECRET"), tool(1, finish_reason="tool_calls"))
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks)
delivered: Final = tuple(
(bool(choice.delta.content), choice.finish_reason is not None)
for chunk in out
for choice in chunk.choices
)
carries_text: Final = tuple(position for position, (text, _) in enumerate(delivered) if text)
finishes: Final = tuple(position for position, (_, finish) in enumerate(delivered) if finish)
assert carries_text and finishes
assert max(carries_text) <= min(finishes)
assert "".join(_delta_text(chunk) for chunk in out) == "checked:hello MASKED"
@pytest.mark.asyncio
async def test_a_mixed_text_and_tool_call_chunk_never_finishes_ahead_of_its_text(self):
"""A client that stops at finish_reason must not lose rewritten text held behind it."""
mixed: Final = ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(
content="the code is SECRET",
role="assistant",
tool_calls=[
{
"index": 0,
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": '{"q": "SECRET"}'},
}
],
),
finish_reason="tool_calls",
)
],
)
leading: Final = _stream_chunk("the code is ")
out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), [leading, mixed])
delivered: Final = "".join(_delta_text(item) for item in out)
assert "MASKED" in delivered
finished_at: Final = next(
(
position
for position, item in enumerate(out)
if any(getattr(choice, "finish_reason", None) for choice in getattr(item, "choices", ()) or ())
),
None,
)
text_at: Final = next(
(position for position, item in enumerate(out) if "MASKED" in _delta_text(item)),
None,
)
assert finished_at is not None
assert text_at is not None
assert finished_at >= text_at
@pytest.mark.asyncio
async def test_incremental_diff_emits_uppercased_deltas(self):
"""incremental_diff: the client receives uppercased deltas whose
@ -1439,14 +1748,15 @@ class TestStreamingTransform:
assert out[-2].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
async def test_tool_call_blocking_guardrail_is_enforced(self):
@pytest.mark.parametrize(("content", "allowed_scans"), [(None, 0), ("proposal", 0), ("proposal", 1)])
async def test_tool_call_blocking_guardrail_is_enforced(self, content: str | None, allowed_scans: int):
"""A guardrail that blocks on tool calls must terminate the incremental_diff
stream: tool calls go through the block decision, not bypass it."""
from litellm.exceptions import GuardrailRaisedException
class _ToolCallBlocker(_StreamingTextGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
if input_type == "response" and inputs.get("tool_calls"):
if input_type == "response" and self.response_calls >= allowed_scans:
raise GuardrailRaisedException(
guardrail_name="tc-block",
message="blocked tool call",
@ -1459,7 +1769,7 @@ class TestStreamingTransform:
StreamingChoices(
index=0,
delta=Delta(
content=None,
content=content,
tool_calls=[
{
"index": 0,
@ -1474,8 +1784,21 @@ class TestStreamingTransform:
],
)
async def upstream():
yield tool_chunk
stream: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"),
response=upstream(),
request_data={"guardrail_to_apply": _ToolCallBlocker(), "model": "gpt-4"},
)
async def consume_checked_stream() -> None:
async for chunk in stream:
assert isinstance(chunk, ModelResponseStream)
assert all(not choice.delta.tool_calls and choice.finish_reason is None for choice in chunk.choices)
with pytest.raises(GuardrailRaisedException):
await _drive_stream(UnifiedLLMGuardrails(), _ToolCallBlocker(), [tool_chunk])
await consume_checked_stream()
@pytest.mark.asyncio
async def test_mixed_content_and_tool_call_chunk_does_not_leak_text(self):
@ -2715,7 +3038,7 @@ class TestStreamingClientDisconnectScan:
yield _stream_chunk(" tail", finish_reason="stop")
@pytest.mark.asyncio
async def test_closing_while_an_incremental_diff_block_is_delivered_does_not_inspect_the_tool_call_again(self):
async def test_incremental_diff_block_withholds_tool_calls_and_does_not_inspect_again(self):
guardrail = _MarkerBlockingStreamingTextGuardrail()
async def upstream() -> AsyncIterator[ModelResponseStream]:
@ -2729,12 +3052,16 @@ class TestStreamingClientDisconnectScan:
released = [await stream.__anext__(), await stream.__anext__()]
await stream.aclose()
assert [call.function.name for call in released[0].choices[0].delta.tool_calls] == ["get_weather"]
assert guardrail.received_texts == [["BLOCKME"]], guardrail.received_texts
assert guardrail.received_tool_calls == [], guardrail.received_tool_calls
assert all(isinstance(item, bytes) for item in released)
assert b"get_weather" not in b"".join(released)
assert b"Paris" not in b"".join(released)
assert guardrail.received_texts == [["BLOCKME tail"]], guardrail.received_texts
assert [[call["function"]["name"] for call in calls] for calls in guardrail.received_tool_calls] == [
["get_weather"]
]
@pytest.mark.asyncio
async def test_closing_after_a_released_tool_call_under_incremental_diff_scans_no_held_back_text(self):
async def test_closing_after_a_buffered_tool_call_inspects_held_text_before_delivery(self):
guardrail = _StreamingTextGuardrail(holdback_schedule=[len("held secret")] * 2)
async def upstream() -> AsyncIterator[ModelResponseStream]:
@ -2743,12 +3070,14 @@ class TestStreamingClientDisconnectScan:
yield tool_chunk
stream = self._guarded_stream(guardrail, upstream())
received = await stream.__anext__()
received: Final = tuple([await stream.__anext__() for _ in range(2)])
scans_before_close: Final = len(guardrail.received_texts)
await stream.aclose()
assert [call.function.name for call in received.choices[0].delta.tool_calls] == ["get_weather"]
assert [call.function.name for call in received[-1].choices[0].delta.tool_calls] == ["get_weather"]
assert guardrail.received_tool_calls, guardrail.received_texts
assert all("held secret" not in text for text in guardrail.received_texts[-1]), guardrail.received_texts
assert guardrail.received_texts[0] == ["held secret"]
assert len(guardrail.received_texts) == scans_before_close
def _responses_delta(sequence_number, text):