mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
refactor(guardrails): separate shared fixes from native integration
This commit is contained in:
albertbausili
parent
42b69e7319
commit
a0c4d90a19
9 changed files with 328 additions and 606 deletions
6
.github/workflows/test-rust.yml
vendored
6
.github/workflows/test-rust.yml
vendored
|
|
@ -185,11 +185,7 @@ jobs:
|
|||
env:
|
||||
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
|
||||
- name: Test native routes from the installed wheel
|
||||
run: |
|
||||
wheel=(dist/*.whl)
|
||||
uv run --isolated --no-project --with "${wheel[0]}" \
|
||||
python tests/unit/rust_bridge/native_route_wheel_test.py "${wheel[0]}"
|
||||
- run: python tests/unit/rust_bridge/native_route_wheel_test.py dist/*.whl
|
||||
|
||||
- name: Run pytest tests/test_litellm_rust with the compiled extension
|
||||
run: make test-rust-extension
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
|
||||
|
||||
|
|
@ -31,12 +31,8 @@ class StreamTransformSink:
|
|||
single call, not shared state.
|
||||
"""
|
||||
|
||||
mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict)
|
||||
holdback_per_choice: Mapping[int, int] = field(default_factory=dict)
|
||||
|
||||
def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None:
|
||||
self.mutated_text_per_choice = texts
|
||||
self.holdback_per_choice = holdback
|
||||
mutated_text_per_choice: dict[int, str] = field(default_factory=dict)
|
||||
holdback_per_choice: dict[int, int] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
|
|
@ -18,11 +18,9 @@ import json
|
|||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from itertools import chain
|
||||
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
|
||||
|
|
@ -48,7 +46,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
stream_item_items,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
coerce_stream_holdback_value,
|
||||
)
|
||||
|
|
@ -170,8 +168,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"] = (
|
||||
|
|
@ -260,12 +258,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
|
||||
names: Final[list[str]] = []
|
||||
for tool in data.get("tools") or ():
|
||||
for tool in data.get("tools") or []:
|
||||
if isinstance(tool, dict) and tool.get("type") == "function":
|
||||
fn = tool.get("function")
|
||||
if isinstance(fn, dict) and fn.get("name"):
|
||||
names.append(str(fn["name"]))
|
||||
for fn in data.get("functions") or ():
|
||||
for fn in data.get("functions") or []:
|
||||
if isinstance(fn, dict) and fn.get("name"):
|
||||
names.append(str(fn["name"]))
|
||||
return names
|
||||
|
|
@ -389,7 +387,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict[str, object] | None = None,
|
||||
request_data: dict | None = None,
|
||||
) -> ModelResponse:
|
||||
"""
|
||||
Process output response by applying guardrails to text content.
|
||||
|
|
@ -440,7 +438,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
# create a standalone dict (SDK / direct-call path).
|
||||
if request_data is None:
|
||||
request_data = {"response": response}
|
||||
request_data.setdefault("response", response)
|
||||
else:
|
||||
if "response" not in request_data:
|
||||
request_data["response"] = response
|
||||
|
||||
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
|
||||
|
||||
|
|
@ -460,7 +460,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)
|
||||
|
|
@ -494,7 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict[str, object] | None = None,
|
||||
request_data: dict | None = None,
|
||||
stream_transform_sink: StreamTransformSink | None = None,
|
||||
deliver_ended_stream_rewrites: bool = False,
|
||||
) -> list["ModelResponseStream"]:
|
||||
|
|
@ -550,14 +550,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
user_api_key_dict: "UserAPIKeyAuth | None",
|
||||
request_data: dict[str, object] | None,
|
||||
request_data: dict | 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 = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far)
|
||||
has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
|
||||
|
||||
if has_stream_ended:
|
||||
await self._process_ended_stream(
|
||||
|
|
@ -623,7 +623,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
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
|
||||
|
|
@ -656,36 +656,13 @@ 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)
|
||||
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,)
|
||||
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_request_data: Final = request_data if request_data is not None else {}
|
||||
try:
|
||||
for inspection_response in inspection_responses:
|
||||
inspection_request_data["response"] = inspection_response
|
||||
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
|
||||
)
|
||||
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,
|
||||
)
|
||||
finally:
|
||||
inspection_request_data["response"] = model_response
|
||||
inspection_request_data["responses"] = responses_so_far
|
||||
if not deliver_ended_stream_rewrites:
|
||||
return
|
||||
await self._write_ended_stream_text_rewrites(
|
||||
|
|
@ -786,7 +763,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
user_api_key_dict: "UserAPIKeyAuth | None",
|
||||
request_data: dict[str, object] | None,
|
||||
request_data: dict | None,
|
||||
sink: StreamTransformSink,
|
||||
) -> None:
|
||||
"""Run the guardrail over the raw accumulated text and report the
|
||||
|
|
@ -798,7 +775,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
"""
|
||||
raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far)
|
||||
if not raw_by_index:
|
||||
sink.record(MappingProxyType({}), MappingProxyType({}))
|
||||
sink.mutated_text_per_choice = {}
|
||||
sink.holdback_per_choice = {}
|
||||
return
|
||||
|
||||
# Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits
|
||||
|
|
@ -810,72 +788,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
indices: Final = sorted(raw_by_index.keys())
|
||||
texts_to_check: Final = [raw_by_index[i] for i in indices]
|
||||
|
||||
resolved_request_data: Final[dict[str, object]] = (
|
||||
request_data if request_data is not None else {"responses": responses_so_far}
|
||||
)
|
||||
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)
|
||||
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 self._streamed_tool_call_fingerprints(responses_so_far):
|
||||
assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj)
|
||||
resolved_request_data["response"] = assembled
|
||||
if len(assembled.choices) > 1:
|
||||
choice_rounds: Final = tuple(
|
||||
(
|
||||
choice,
|
||||
StreamTransformSink(),
|
||||
[self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far],
|
||||
)
|
||||
for choice in assembled.choices
|
||||
)
|
||||
try:
|
||||
for choice, choice_sink, choice_chunks in choice_rounds:
|
||||
resolved_request_data["response"] = assembled.model_copy(
|
||||
update=MappingProxyType({"choices": [choice]})
|
||||
)
|
||||
resolved_request_data["responses"] = choice_chunks
|
||||
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,
|
||||
)
|
||||
finally:
|
||||
resolved_request_data["response"] = assembled
|
||||
resolved_request_data["responses"] = responses_so_far
|
||||
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
|
||||
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,
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
|
@ -895,14 +819,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
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)
|
||||
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)
|
||||
}
|
||||
|
||||
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))
|
||||
|
|
@ -1222,18 +1145,20 @@ 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(
|
||||
(choice.index, tool_call)
|
||||
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((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)
|
||||
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)
|
||||
return tuple(
|
||||
tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key)
|
||||
for key in sorted(identified & named)
|
||||
tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
|
||||
)
|
||||
|
||||
def _write_ended_stream_tool_call_rewrites(
|
||||
|
|
@ -1243,10 +1168,28 @@ 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
|
||||
|
||||
|
|
@ -1269,7 +1212,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
async def _apply_guardrail_responses_to_output_streaming(
|
||||
self,
|
||||
responses: list["ModelResponseStream"],
|
||||
guardrailed_texts: Sequence[str],
|
||||
guardrailed_texts: list[str],
|
||||
task_mappings: list[tuple[int, int | None]],
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -1287,9 +1230,12 @@ 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 = MappingProxyType(
|
||||
{task_mappings[index]: text for index, text in enumerate(guardrailed_texts)}
|
||||
)
|
||||
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
|
||||
|
||||
# Track which choices we've already set the guardrailed text for
|
||||
# Key: (choice_idx, content_idx), Value: boolean (True if already set)
|
||||
|
|
@ -1299,20 +1245,34 @@ 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, 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
|
||||
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
|
||||
|
||||
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:
|
||||
message.content = guardrail_map[str_key] if str_key not in already_set else ""
|
||||
already_set[str_key] = True
|
||||
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 = ""
|
||||
|
||||
elif isinstance(content, list):
|
||||
# List content - handle each content item
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union,
|
|||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
|
|
@ -2318,26 +2318,6 @@ async def test_custom_code_guardrail(
|
|||
)
|
||||
|
||||
|
||||
_GUARDRAIL_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _metadata_fields(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _GUARDRAIL_METADATA_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return {}
|
||||
|
||||
|
||||
def _guardrail_request_metadata(caller: object, proxy: object) -> Mapping[str, object]:
|
||||
caller_fields: Final = tuple(
|
||||
(key, value) for key, value in _metadata_fields(caller).items() if not key.startswith("user_api_key_")
|
||||
)
|
||||
key_identity: Final = tuple(
|
||||
(key, value) for key, value in _metadata_fields(proxy).items() if key.startswith("user_api_key_")
|
||||
)
|
||||
return dict((*caller_fields, *key_identity))
|
||||
|
||||
|
||||
def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
|
|
@ -2517,14 +2497,9 @@ async def apply_guardrail(
|
|||
if litellm_logging_obj is not None:
|
||||
_patch_logging_obj_for_guardrail(litellm_logging_obj, request)
|
||||
|
||||
metadata: Final = _guardrail_request_metadata(request.metadata, data.get("metadata"))
|
||||
request_data: Final[dict] = {
|
||||
key: value
|
||||
for key, value in (
|
||||
("messages", request.messages),
|
||||
("metadata", metadata if request.metadata is not None or metadata else None),
|
||||
)
|
||||
if value is not None
|
||||
**({"messages": request.messages} if request.messages is not None else {}),
|
||||
**({"metadata": request.metadata} if request.metadata is not None else {}),
|
||||
}
|
||||
_input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type)
|
||||
guardrailed_inputs: Final = await active_guardrail.apply_guardrail(
|
||||
|
|
|
|||
|
|
@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime
|
|||
|
||||
import anyio
|
||||
from fastapi import HTTPException
|
||||
from pydantic import InstanceOf, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
|
|
@ -50,7 +49,6 @@ 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)
|
||||
|
||||
_RequestData: TypeAlias = dict[str, object]
|
||||
|
||||
|
|
@ -515,9 +513,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
self,
|
||||
*,
|
||||
reference_chunk: object,
|
||||
mutated_text_per_choice: Mapping[int, str],
|
||||
mutated_text_per_choice: dict[int, str],
|
||||
emitted_text_per_choice: dict[int, str],
|
||||
holdback_per_choice: Mapping[int, int],
|
||||
holdback_per_choice: dict[int, int],
|
||||
finish_reason_per_choice: dict[int, str | None],
|
||||
held_chars_per_choice: dict[int, int],
|
||||
is_final: bool,
|
||||
|
|
@ -753,13 +751,46 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
|
||||
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
|
||||
|
||||
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):
|
||||
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):
|
||||
|
|
@ -777,62 +808,39 @@ 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:
|
||||
inspected_responses: Final = copy.deepcopy(responses_so_far)
|
||||
async for out in self._inspect_full_response(
|
||||
async for out in self._inspect_full_response_for_block(
|
||||
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=inspected_responses,
|
||||
responses_so_far=copy.deepcopy(responses_so_far),
|
||||
responses_yielded=responses_yielded,
|
||||
):
|
||||
yield out
|
||||
|
||||
if saw_text_content:
|
||||
async for out in _round(last_chunk, is_final=False):
|
||||
yield out
|
||||
tool_chunks: Final = tuple(
|
||||
self._tool_call_passthrough_chunk(
|
||||
buffered_item,
|
||||
finish_reason_per_choice=finish_reason_per_choice,
|
||||
held_choices=_held_choices(held_chars_per_choice),
|
||||
)
|
||||
for buffered_item in inspected_responses
|
||||
if self._chunk_has_tool_calls(buffered_item)
|
||||
)
|
||||
|
||||
async def checked_tail() -> AsyncGenerator[object, None]:
|
||||
try:
|
||||
async for tail_chunk in self._emit_stream_tail(
|
||||
last_chunk=last_chunk,
|
||||
final_round=_round,
|
||||
responses_so_far=responses_so_far,
|
||||
responses_yielded=responses_yielded,
|
||||
):
|
||||
yield tail_chunk
|
||||
except _StreamTerminated as exc:
|
||||
yield exc
|
||||
|
||||
tail_chunks: Final = tuple([chunk async for chunk in checked_tail()])
|
||||
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,
|
||||
|
|
@ -842,6 +850,37 @@ 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(
|
||||
*,
|
||||
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,
|
||||
)
|
||||
|
||||
async def _emit_stream_tail(
|
||||
self,
|
||||
|
|
@ -860,21 +899,29 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
responses_yielded.append(trailing)
|
||||
yield trailing
|
||||
|
||||
async def _inspect_full_response(
|
||||
async def _inspect_full_response_for_block(
|
||||
self,
|
||||
*,
|
||||
endpoint_translation: _EndpointTranslation,
|
||||
guardrail_to_apply: CustomGuardrail,
|
||||
request_data: _RequestData,
|
||||
request_data: dict,
|
||||
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=True):
|
||||
with anyio.CancelScope(shield=bool(responses_yielded)):
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
|
|
@ -882,10 +929,7 @@ 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
|
||||
|
|
@ -1069,15 +1113,12 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
# litellm.integrations.custom_guardrail.
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
|
||||
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))
|
||||
)
|
||||
if guardrail_to_apply is None:
|
||||
guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
|
||||
typed_request_data: Final[_RequestData] = request_data
|
||||
|
||||
def _streaming_flag(name: str, default: object) -> Any:
|
||||
return self.resolve_streaming_flag(resolved_guardrail, name, default)
|
||||
return self.resolve_streaming_flag(guardrail_to_apply, name, default)
|
||||
|
||||
sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
|
||||
# Only apply the guardrail at end of stream (not per chunk).
|
||||
|
|
@ -1100,30 +1141,30 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
|
||||
if (
|
||||
buffer_until_moderated
|
||||
and resolved_guardrail is not None
|
||||
and getattr(resolved_guardrail, "mask_response_content", False)
|
||||
and guardrail_to_apply is not None
|
||||
and getattr(guardrail_to_apply, "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.",
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
buffer_until_moderated = False
|
||||
|
||||
if buffer_until_moderated and not release_on_scan:
|
||||
end_of_stream_only = True
|
||||
|
||||
if resolved_guardrail is None:
|
||||
if guardrail_to_apply is None:
|
||||
async for item in response:
|
||||
yield item
|
||||
return
|
||||
|
||||
event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call
|
||||
if resolved_guardrail.should_run_guardrail(data=request_data, event_type=event_type) is not True:
|
||||
if guardrail_to_apply.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",
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
async for item in response:
|
||||
yield item
|
||||
|
|
@ -1143,7 +1184,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
if transform_call_type is not None:
|
||||
async with contextlib.aclosing(
|
||||
self._run_incremental_transform_stream(
|
||||
guardrail_to_apply=resolved_guardrail,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
response=response,
|
||||
request_data=typed_request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -1160,7 +1201,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(resolved_guardrail, "guardrail_name", None),
|
||||
getattr(guardrail_to_apply, "guardrail_name", None),
|
||||
)
|
||||
|
||||
# Infer call type from first chunk
|
||||
|
|
@ -1232,7 +1273,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,
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
if buffer_until_moderated:
|
||||
if hold_window:
|
||||
|
|
@ -1252,7 +1293,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
"Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
|
||||
chunk_counter,
|
||||
sampling_rate,
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
original_items = (
|
||||
|
|
@ -1262,7 +1303,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
try:
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=resolved_guardrail,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
|
|
@ -1308,7 +1349,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),
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
withheld_items[:] = original_items
|
||||
continue
|
||||
|
|
@ -1328,7 +1369,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
verbose_proxy_logger.debug(
|
||||
"Processing final streaming response with all %s chunks for guardrail %s",
|
||||
len(responses_so_far),
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
|
||||
endpoint_translation = mappings[CallTypes(call_type)]()
|
||||
|
|
@ -1345,7 +1386,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
if _is_redundant_scan(end_scan_key, last_scan_key):
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
|
||||
resolved_guardrail.guardrail_name,
|
||||
guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
for buffered_item in buffered_items or ():
|
||||
yield buffered_item
|
||||
|
|
@ -1358,7 +1399,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
with anyio.CancelScope(shield=chunks_yielded):
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=resolved_guardrail,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
|
|
@ -1402,13 +1443,13 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
chunks_yielded
|
||||
and not verdict_settled
|
||||
and translation_class is not None
|
||||
and isinstance(resolved_guardrail, CustomGuardrail)
|
||||
and isinstance(guardrail_to_apply, CustomGuardrail)
|
||||
):
|
||||
await self._scan_released_stream_after_disconnect(
|
||||
endpoint_translation=translation_class(),
|
||||
responses_released=responses_yielded,
|
||||
last_scan_key=last_scan_key,
|
||||
guardrail_to_apply=resolved_guardrail,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=typed_request_data,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10,18 +10,11 @@ Run with: pytest tests/unit/interactions/test_openapi_compliance.py -v
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Dict, Final
|
||||
from typing import Any, Dict
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from jsonschema import Draft202012Validator
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig
|
||||
from litellm.types.interactions import InteractionInput
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from openapi_core import OpenAPI
|
||||
|
||||
OPENAPI_SPEC_URL = "https://ai.google.dev/static/api/interactions.openapi.json"
|
||||
|
|
@ -45,6 +38,25 @@ def _load_openapi_spec_dict() -> Dict[str, Any]:
|
|||
)
|
||||
|
||||
|
||||
def _model_create_request_schema(spec_dict: Dict[str, Any]) -> Dict[str, Any]:
|
||||
schemas = spec_dict["components"]["schemas"]
|
||||
create_path = next(path for path in spec_dict["paths"] if path.endswith("/interactions"))
|
||||
body_schema = spec_dict["paths"][create_path]["post"]["requestBody"]["content"]["application/json"]["schema"]
|
||||
variants = [schemas[option["$ref"].split("/")[-1]] for option in body_schema.get("oneOf", []) if "$ref" in option]
|
||||
return next(variant for variant in variants if "model" in variant.get("properties", {}))
|
||||
|
||||
|
||||
def _interaction_resource_path(spec_dict: Dict[str, Any], method: str) -> str | None:
|
||||
return next(
|
||||
(
|
||||
path
|
||||
for path, methods in spec_dict["paths"].items()
|
||||
if re.search(r"/interactions/\{[^}]+\}$", path) and method in methods
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _declared_type_value(variant_schema: Dict[str, Any]) -> Any:
|
||||
"""The single `type` value a union variant pins, whether spelled as a const or a 1-item enum."""
|
||||
type_property = variant_schema.get("properties", {}).get("type", {})
|
||||
|
|
@ -64,58 +76,59 @@ def openapi_spec(spec_dict: Dict[str, Any]) -> OpenAPI:
|
|||
return OpenAPI.from_dict(spec_dict)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_request_schema(spec_dict: Mapping[str, object]) -> Mapping[str, object]:
|
||||
objects: Final = TypeAdapter(Mapping[str, Mapping[str, object]])
|
||||
components: Final = TypeAdapter(Mapping[str, object]).validate_python(spec_dict["components"])
|
||||
schemas: Final = objects.validate_python(components["schemas"])
|
||||
paths: Final = objects.validate_python(spec_dict["paths"])
|
||||
operation: Final = next(methods["post"] for path, methods in paths.items() if path.endswith("/interactions"))
|
||||
post: Final = TypeAdapter(Mapping[str, object]).validate_python(operation)
|
||||
body: Final = TypeAdapter(Mapping[str, object]).validate_python(post["requestBody"])
|
||||
content: Final = objects.validate_python(body["content"])
|
||||
schema: Final = TypeAdapter(Mapping[str, object]).validate_python(content["application/json"]["schema"])
|
||||
variants: Final = TypeAdapter(tuple[Mapping[str, str], ...]).validate_python(schema["oneOf"])
|
||||
return next(
|
||||
schemas[variant["$ref"].rsplit("/", 1)[-1]]
|
||||
for variant in variants
|
||||
if "model" in TypeAdapter(Mapping[str, object]).validate_python(
|
||||
schemas[variant["$ref"].rsplit("/", 1)[-1]]["properties"]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TestRequestCompliance:
|
||||
def test_create_model_interaction_request_schema(self, model_request_schema: Mapping[str, object]) -> None:
|
||||
config: Final = GoogleAIStudioInteractionsConfig()
|
||||
payload: Final = config.transform_request(
|
||||
model="test-model", agent=None, input="test input",
|
||||
optional_params={"stream": True, "response_mime_type": "application/json"},
|
||||
litellm_params=GenericLiteLLMParams(api_key="test-key"), headers={},
|
||||
)
|
||||
properties: Final = TypeAdapter(Mapping[str, object]).validate_python(model_request_schema["properties"])
|
||||
assert payload == {
|
||||
"model": "test-model", "input": "test input", "stream": True,
|
||||
"response_format": {"type": "text", "mime_type": "application/json"},
|
||||
}
|
||||
assert payload.keys() <= properties.keys()
|
||||
assert set(config.get_supported_params("test-model")) - {"agent", "response_mime_type"} <= properties.keys()
|
||||
"""Tests that our request bodies match the OpenAPI spec."""
|
||||
|
||||
@pytest.mark.parametrize("input_value", ["test input", [{"type": "text", "text": "test input"}]])
|
||||
def test_input_types_match_spec(
|
||||
self, spec_dict: Mapping[str, object], model_request_schema: Mapping[str, object],
|
||||
input_value: InteractionInput,
|
||||
) -> None:
|
||||
payload: Final = GoogleAIStudioInteractionsConfig().transform_request(
|
||||
model="test-model", agent=None, input=input_value, optional_params={},
|
||||
litellm_params=GenericLiteLLMParams(api_key="test-key"), headers={},
|
||||
)
|
||||
properties: Final = TypeAdapter(Mapping[str, Mapping[str, object]]).validate_python(
|
||||
model_request_schema["properties"]
|
||||
)
|
||||
validator: Final = Draft202012Validator(spec_dict).evolve(schema=properties["input"])
|
||||
assert not tuple(validator.iter_errors(payload["input"]))
|
||||
assert payload["input"] == input_value
|
||||
def test_create_model_interaction_request_schema(self, spec_dict):
|
||||
schema = _model_create_request_schema(spec_dict)
|
||||
|
||||
assert "model" in schema["required"]
|
||||
assert "input" in schema["properties"]
|
||||
|
||||
# Check our supported optional fields exist in spec
|
||||
our_optional_fields = [
|
||||
"tools",
|
||||
"system_instruction",
|
||||
"generation_config",
|
||||
"stream",
|
||||
"store",
|
||||
"background",
|
||||
"response_modalities",
|
||||
"response_format",
|
||||
"response_mime_type",
|
||||
"previous_interaction_id",
|
||||
]
|
||||
|
||||
spec_properties = schema["properties"]
|
||||
for field in our_optional_fields:
|
||||
assert field in spec_properties, f"Field '{field}' not in OpenAPI spec"
|
||||
print(f"✓ Field '{field}' exists in spec")
|
||||
|
||||
def test_input_types_match_spec(self, spec_dict):
|
||||
"""Verify input field supports string, Content, Content[], Turn[]."""
|
||||
schema = _model_create_request_schema(spec_dict)
|
||||
input_schema = schema["properties"]["input"]
|
||||
|
||||
# The input property may be inline oneOf or a $ref to InteractionsInput
|
||||
if "$ref" in input_schema:
|
||||
ref_name = input_schema["$ref"].split("/")[-1]
|
||||
input_schema = spec_dict["components"]["schemas"][ref_name]
|
||||
|
||||
# Should be oneOf with multiple types
|
||||
assert "oneOf" in input_schema
|
||||
|
||||
input_types = []
|
||||
for option in input_schema["oneOf"]:
|
||||
if option.get("type") == "string":
|
||||
input_types.append("string")
|
||||
elif option.get("type") == "array":
|
||||
input_types.append("array")
|
||||
elif "$ref" in option:
|
||||
input_types.append(option["$ref"])
|
||||
|
||||
print(f"Input supports types: {input_types}")
|
||||
assert "string" in input_types, "Input should support string"
|
||||
assert "array" in input_types, "Input should support array"
|
||||
|
||||
def test_content_variants_are_identified_by_their_type_field(self, spec_dict):
|
||||
"""Verify a Content part can be told apart by its `type`, however the spec spells that.
|
||||
|
|
@ -312,24 +325,19 @@ class TestEndpointCompliance:
|
|||
assert create_path is not None, "POST /interactions endpoint not found"
|
||||
print(f"✓ Create endpoint: POST {create_path}")
|
||||
|
||||
@pytest.mark.parametrize("method", ["get", "delete"])
|
||||
def test_interaction_item_url_matches_spec(self, spec_dict: Mapping[str, object], method: str) -> None:
|
||||
config: Final = GoogleAIStudioInteractionsConfig()
|
||||
transform: Final = (
|
||||
config.transform_get_interaction_request if method == "get" else config.transform_delete_interaction_request
|
||||
)
|
||||
url, body = transform(
|
||||
interaction_id="test-interaction", api_base="https://example.com",
|
||||
litellm_params=GenericLiteLLMParams(api_key="test-key"), headers={},
|
||||
)
|
||||
path: Final = httpx.URL(url).path
|
||||
paths: Final = TypeAdapter(Mapping[str, Mapping[str, object]]).validate_python(spec_dict["paths"])
|
||||
assert any(
|
||||
re.fullmatch(re.sub(r"\{[^}]+\}", "[^/]+", template), path) and method in operations
|
||||
for template, operations in paths.items()
|
||||
), path
|
||||
assert path.endswith("/test-interaction")
|
||||
assert body == {}
|
||||
def test_get_endpoint_exists(self, spec_dict):
|
||||
"""Verify GET /interactions/{id} endpoint exists."""
|
||||
get_path = _interaction_resource_path(spec_dict, "get")
|
||||
|
||||
assert get_path is not None, "GET /interactions/{id} endpoint not found"
|
||||
print(f"✓ Get endpoint: GET {get_path}")
|
||||
|
||||
def test_delete_endpoint_exists(self, spec_dict):
|
||||
"""Verify DELETE /interactions/{id} endpoint exists."""
|
||||
delete_path = _interaction_resource_path(spec_dict, "delete")
|
||||
|
||||
assert delete_path is not None, "DELETE /interactions/{id} endpoint not found"
|
||||
print(f"✓ Delete endpoint: DELETE {delete_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ with guardrail transformations, including tool calls.
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final, Literal, Optional
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -1395,65 +1395,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
|
|||
)
|
||||
|
||||
assert [
|
||||
[(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"}')]]
|
||||
(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"}')]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@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
|
||||
async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
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
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
chunks = self._two_choice_tool_call_stream_chunks()
|
||||
|
||||
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,
|
||||
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,
|
||||
)
|
||||
assert request_data["visited_choices"] == (0, 1)
|
||||
assert request_data["metadata"]["trace"] == "retained"
|
||||
assert request_data["responses"] is chunks
|
||||
restored: Final = request_data["response"]
|
||||
assert isinstance(restored, ModelResponse)
|
||||
assert tuple(choice.index for choice in restored.choices) == (0, 1)
|
||||
|
||||
@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,
|
||||
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"
|
||||
)
|
||||
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):
|
||||
|
|
|
|||
|
|
@ -8,8 +8,6 @@ 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
|
||||
|
||||
|
|
@ -52,7 +50,6 @@ from litellm.types.llms.openai import ResponsesAPIResponse
|
|||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
CallTypesLiteral,
|
||||
ChatCompletionMessageToolCall,
|
||||
Delta,
|
||||
GenericGuardrailAPIInputs,
|
||||
ModelResponse,
|
||||
|
|
@ -1011,42 +1008,6 @@ 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."""
|
||||
|
|
@ -1055,103 +1016,6 @@ 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
|
||||
|
|
@ -1572,15 +1436,14 @@ class TestStreamingTransform:
|
|||
assert out[-2].choices[0].finish_reason == "stop"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@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):
|
||||
async def test_tool_call_blocking_guardrail_is_enforced(self):
|
||||
"""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 self.response_calls >= allowed_scans:
|
||||
if input_type == "response" and inputs.get("tool_calls"):
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name="tc-block",
|
||||
message="blocked tool call",
|
||||
|
|
@ -1593,7 +1456,7 @@ class TestStreamingTransform:
|
|||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
content=None,
|
||||
tool_calls=[
|
||||
{
|
||||
"index": 0,
|
||||
|
|
@ -1608,21 +1471,8 @@ 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 consume_checked_stream()
|
||||
await _drive_stream(UnifiedLLMGuardrails(), _ToolCallBlocker(), [tool_chunk])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_content_and_tool_call_chunk_does_not_leak_text(self):
|
||||
|
|
@ -2840,7 +2690,7 @@ class TestStreamingClientDisconnectScan:
|
|||
yield _stream_chunk(" tail", finish_reason="stop")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_incremental_diff_block_withholds_tool_calls_and_does_not_inspect_again(self):
|
||||
async def test_closing_while_an_incremental_diff_block_is_delivered_does_not_inspect_the_tool_call_again(self):
|
||||
guardrail = _MarkerBlockingStreamingTextGuardrail()
|
||||
|
||||
async def upstream() -> AsyncIterator[ModelResponseStream]:
|
||||
|
|
@ -2854,16 +2704,12 @@ class TestStreamingClientDisconnectScan:
|
|||
released = [await stream.__anext__(), await stream.__anext__()]
|
||||
await stream.aclose()
|
||||
|
||||
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"]
|
||||
]
|
||||
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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closing_after_a_buffered_tool_call_inspects_held_text_before_delivery(self):
|
||||
async def test_closing_after_a_released_tool_call_under_incremental_diff_scans_no_held_back_text(self):
|
||||
guardrail = _StreamingTextGuardrail(holdback_schedule=[len("held secret")] * 2)
|
||||
|
||||
async def upstream() -> AsyncIterator[ModelResponseStream]:
|
||||
|
|
@ -2872,14 +2718,12 @@ class TestStreamingClientDisconnectScan:
|
|||
yield tool_chunk
|
||||
|
||||
stream = self._guarded_stream(guardrail, upstream())
|
||||
received: Final = tuple([await stream.__anext__() for _ in range(2)])
|
||||
scans_before_close: Final = len(guardrail.received_texts)
|
||||
received = await stream.__anext__()
|
||||
await stream.aclose()
|
||||
|
||||
assert [call.function.name for call in received[-1].choices[0].delta.tool_calls] == ["get_weather"]
|
||||
assert [call.function.name for call in received.choices[0].delta.tool_calls] == ["get_weather"]
|
||||
assert guardrail.received_tool_calls, guardrail.received_texts
|
||||
assert guardrail.received_texts[0] == ["held secret"]
|
||||
assert len(guardrail.received_texts) == scans_before_close
|
||||
assert all("held secret" not in text for text in guardrail.received_texts[-1]), guardrail.received_texts
|
||||
|
||||
|
||||
def _responses_delta(sequence_number, text):
|
||||
|
|
|
|||
|
|
@ -2127,7 +2127,7 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
|
|||
}
|
||||
|
||||
|
||||
def _patch_apply_guardrail_env(mocker, guardrail_result, processed_data=None):
|
||||
def _patch_apply_guardrail_env(mocker, guardrail_result):
|
||||
mock_guardrail = mocker.Mock()
|
||||
mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result)
|
||||
|
||||
|
|
@ -2140,7 +2140,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result, processed_data=None):
|
|||
mock_logging_obj.model_call_details = {}
|
||||
mock_processor = mocker.Mock()
|
||||
mock_processor.common_processing_pre_call_logic = AsyncMock(
|
||||
return_value=(processed_data or {"guardrail_name": "test-guardrail"}, mock_logging_obj)
|
||||
return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
|
||||
|
|
@ -2211,68 +2211,6 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key(mocker):
|
||||
"""Identity fields come from the proxy's own sanitized metadata, never from
|
||||
the caller, so a request cannot impersonate another key or probe its policy."""
|
||||
mock_guardrail = _patch_apply_guardrail_env(
|
||||
mocker,
|
||||
{"texts": ["ok"]},
|
||||
processed_data={
|
||||
"guardrail_name": "test-guardrail",
|
||||
"metadata": {
|
||||
"route": "/apply_guardrail",
|
||||
"user_api_key_alias": "billing-app",
|
||||
"user_api_key_user_id": "u-1",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-guardrail",
|
||||
text="hello",
|
||||
metadata={
|
||||
"forbidden_topics": ["tax"],
|
||||
"user_api_key_alias": "someone-else",
|
||||
"user_api_key_user_email": "victim@example.com",
|
||||
},
|
||||
)
|
||||
await apply_guardrail(
|
||||
fastapi_request=mocker.Mock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"),
|
||||
)
|
||||
|
||||
forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"]
|
||||
assert forwarded == {
|
||||
"forbidden_topics": ["tax"],
|
||||
"user_api_key_alias": "billing-app",
|
||||
"user_api_key_user_id": "u-1",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata(mocker):
|
||||
"""A caller that sends no metadata still gets the authenticated identity
|
||||
forwarded, the same shape every LLM route gives a guardrail."""
|
||||
mock_guardrail = _patch_apply_guardrail_env(
|
||||
mocker,
|
||||
{"texts": ["ok"]},
|
||||
processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}},
|
||||
)
|
||||
|
||||
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
|
||||
await apply_guardrail(
|
||||
fastapi_request=mocker.Mock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"),
|
||||
)
|
||||
|
||||
assert mock_guardrail.apply_guardrail.await_args.kwargs["request_data"] == {
|
||||
"metadata": {"user_api_key_alias": "billing-app"}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
|
||||
"""Without metadata, request_data stays empty (backward-compatible)."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue