diff --git a/litellm/litellm_core_utils/served_output_texts.py b/litellm/litellm_core_utils/served_output_texts.py index 854785f304b..b19973e06c3 100644 --- a/litellm/litellm_core_utils/served_output_texts.py +++ b/litellm/litellm_core_utils/served_output_texts.py @@ -4,12 +4,13 @@ post-call guardrail rewrote rather than the provider response the proxy assemble from __future__ import annotations from collections.abc import Sequence +from itertools import chain from typing import Final, Literal -from pydantic import BaseModel, Field, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import ModelResponse, ModelResponseStream +from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts" @@ -44,6 +45,19 @@ class _ChatChoices(LiteLLMBaseModel): choices: list[object] +class _TextCompletionChoice(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + index: int | None = None + text: str | None = None + + +class _TextCompletionChunk(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + choices: tuple[_TextCompletionChoice, ...] + + def _as_json_object(response: object) -> dict[str, object] | None: candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response try: @@ -108,8 +122,13 @@ def served_output_texts(response: object) -> ServedTexts | None: def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None: - if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks): - return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream))) + chat_chunks: Final = tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)) + if chunks and len(chat_chunks) == len(chunks): + return _indexed_texts_by_choice(tuple(chain.from_iterable(_chat_chunk_texts(chunk) for chunk in chat_chunks))) + if chunks and all(isinstance(chunk, TextCompletionResponse) for chunk in chunks): + return _indexed_texts_by_choice( + tuple(chain.from_iterable(_text_completion_chunk_texts(chunk) for chunk in chunks)) + ) from litellm.llms.anthropic.messages.stream_assembly import ( assemble_anthropic_sse_stream, is_anthropic_sse_stream, @@ -121,19 +140,26 @@ def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None: return _chat_texts(assembled) if assembled is not None else None -def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None: - contents: Final = tuple( - content - for chunk in chunks +def _chat_chunk_texts(chunk: ModelResponseStream) -> tuple[tuple[int, str | None], ...]: + return tuple( + (choice.index or 0, content if isinstance(content := choice.delta.content, str) else None) for choice in chunk.choices - if (choice.index or 0) == index and isinstance(content := choice.delta.content, str) ) - return "".join(contents) if contents else None -def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None: - choice_count: Final = max(((choice.index or 0) + 1 for chunk in chunks for choice in chunk.choices), default=0) - texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count)) +def _text_completion_chunk_texts(chunk: object) -> tuple[tuple[int, str | None], ...]: + parsed: Final = _TextCompletionChunk.model_validate(_as_json_object(chunk)) + return tuple((choice.index or 0, choice.text) for choice in parsed.choices) + + +def _joined_choice_text(indexed_texts: Sequence[tuple[int, str | None]], index: int) -> str | None: + texts: Final = tuple(text for text_index, text in indexed_texts if text_index == index and text is not None) + return "".join(texts) if texts else None + + +def _indexed_texts_by_choice(indexed_texts: Sequence[tuple[int, str | None]]) -> ServedTexts | None: + choice_count: Final = max((index + 1 for index, _ in indexed_texts), default=0) + texts: Final = tuple(_joined_choice_text(indexed_texts, index) for index in range(choice_count)) return texts if any(text is not None for text in texts) else None @@ -170,6 +196,8 @@ def _choice_with_text(choice: object, text: str | None) -> object: if choice_obj is None or text is None: return choice message: Final = _as_json_object(choice_obj.get("message")) - if message is None or message.get("content") == text: - return choice - return {**choice_obj, "message": {**message, "content": text}} + if message is not None: + return choice if message.get("content") == text else {**choice_obj, "message": {**message, "content": text}} + if isinstance(choice_obj.get("text"), str): + return {**choice_obj, "text": text} + return choice diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index c4bc04d4627..27c82cfc004 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -5,18 +5,30 @@ This module provides guardrail translation support for OpenAI's text completion The handler processes the 'prompt' parameter for guardrails. """ -from collections.abc import Mapping +import json +from collections.abc import Mapping, Sequence +from itertools import chain from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, + StreamTransformSink, +) +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_chat_stream_usage, + stream_item_field, + stream_item_items, +) +from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from fastapi import HTTPException + + from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.types.utils import TextCompletionResponse class OpenAITextCompletionHandler(BaseTranslation): @@ -32,6 +44,9 @@ class OpenAITextCompletionHandler(BaseTranslation): - A list of strings (for batch completions) """ + delivers_ended_stream_rewrites = True + assembles_streamed_response = True + async def process_input_messages( self, data: dict, @@ -125,7 +140,7 @@ class OpenAITextCompletionHandler(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, ) -> "TextCompletionResponse": """ Process output response by applying guardrails to completion text. @@ -162,11 +177,7 @@ class OpenAITextCompletionHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) # Include model information from the response if available @@ -195,3 +206,168 @@ class OpenAITextCompletionHandler(BaseTranslation): ) return response + + async def process_output_streaming_response( + self, + responses_so_far: list[object], # mutable-ok: rewrites the caller's buffered chunks in place + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: dict[str, object] | None = None, # mutable-ok: the hooks share and mutate the request payload + stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, + ) -> Sequence[object]: + chunks: Final = _text_completion_chunks(responses_so_far) + if not self._check_streaming_has_ended(chunks): + return responses_so_far + assembled: Final = self._assemble_ended_stream(chunks) + pre_guardrail_texts: Final = _choice_texts(assembled) + await self.process_output_response( + response=assembled, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + if not deliver_ended_stream_rewrites: + return responses_so_far + rewrites: Final = tuple( + after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before + ) + for choice_index, text in rewrites: + self._write_choice_text(chunks, choice_index, text) + return responses_so_far + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = _text_completion_chunks(responses_so_far) + if not self._check_streaming_has_ended(chunks): + return None + assembled: Final = self._assemble_ended_stream(chunks) + return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True) + + def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]: + chunks: Final = _text_completion_chunks(responses_so_far) + streamed_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + finished_indices: Final = frozenset( + _choice_index(choice) for choice in streamed_choices if choice.finish_reason + ) + unfinished_indices: Final = tuple( + dict.fromkeys( + _choice_index(choice) for choice in streamed_choices if _choice_index(choice) not in finished_indices + ) + ) + if not unfinished_indices: + return tuple(responses_so_far) + terminator: Final = TextCompletionResponse( + id=chunks[0].id, + model=chunks[0].model, + object="text_completion", + choices=[TextChoices(index=index, text="", finish_reason="length") for index in unfinished_indices], + ) + return (*responses_so_far, terminator) + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + streamed: Final = _text_completion_chunks(responses_so_far or ()) + source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model) + message_chunk: Final = TextCompletionResponse( + id=source.id, + created=source.created, + model=source.model or exc.model, + object="text_completion", + choices=[TextChoices(index=0, text=exc.message)], + ).model_copy(update={"usage": None}) + final_chunk: Final = TextCompletionResponse( + id=source.id, + created=source.created, + model=source.model or exc.model, + object="text_completion", + choices=[TextChoices(index=0, text="", finish_reason="content_filter")], + ).model_copy(update={"usage": _blocked_stream_usage(exc.original_response)}) + return tuple( + f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk) + ) + + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes] | None: + from litellm.proxy.common_request_processing import sse_error_payload + + _, error_obj = sse_error_payload(exc) + return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),) + + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: + streamed_choices: Final = chain.from_iterable(stream_item_items(item, "choices") for item in responses_so_far) + return any(stream_item_field(choice, "finish_reason") is not None for choice in streamed_choices) + + @staticmethod + def _assemble_ended_stream(responses_so_far: Sequence[TextCompletionResponse]) -> TextCompletionResponse: + streamed_choices: Final = tuple(chain.from_iterable(item.choices for item in responses_so_far)) + choice_indices: Final = tuple(dict.fromkeys(_choice_index(choice) for choice in streamed_choices)) + assembled_choices: Final = [ + TextChoices( + index=index, + text="".join(_choice_text(choice) for choice in streamed_choices if _choice_index(choice) == index), + finish_reason=next( + ( + choice.finish_reason + for choice in streamed_choices + if _choice_index(choice) == index and choice.finish_reason + ), + None, + ), + ) + for index in choice_indices + ] + first: Final = responses_so_far[0] + return TextCompletionResponse( + id=first.id, + created=first.created, + model=first.model, + object="text_completion", + choices=assembled_choices, + usage=next((item.usage for item in reversed(responses_so_far) if item.usage is not None), None), + ) + + @staticmethod + def _write_choice_text(responses_so_far: Sequence[TextCompletionResponse], choice_index: int, text: str) -> None: + streamed_choices: Final = chain.from_iterable(item.choices for item in responses_so_far) + head, *tail = (choice for choice in streamed_choices if _choice_index(choice) == choice_index) + head.text = text + for choice in tail: + if isinstance(stream_item_field(choice, "text"), str): + choice.text = "" + + +def _blocked_stream_usage(original_response: object) -> Usage | None: + prompt_tokens, completion_tokens = blocked_chat_stream_usage(original_response) + if isinstance(original_response, list) and prompt_tokens == completion_tokens == 0: + return None + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + +def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]: + return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse)) + + +def _choice_text(choice: TextChoices) -> str: + text: Final = stream_item_field(choice, "text") + return text if isinstance(text, str) else "" + + +def _choice_index(choice: TextChoices) -> int: + return choice.index or 0 + + +def _choice_texts(response: TextCompletionResponse) -> tuple[tuple[int, str], ...]: + return tuple((_choice_index(choice), _choice_text(choice)) for choice in response.choices) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4884e75af4e..7e487ded569 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4059,13 +4059,13 @@ class ProxyBaseLLMRequestProcessing: (assembled_response, cache_hit); the closure also runs non-apply_guardrail post-call hooks via _run_deferred_stream_guardrails. - - Every other anthropic_messages/aresponses stream gets a closure - that dispatches on the stored args shape, because the arming site - cannot tell the producers apart: native iterators store a single - ready-made logging coroutine to enqueue, while bridged streams - (LiteLLMCompletionStreamingIterator, and the plain SSE generator - AnthropicStreamWrapper returns for bridged /v1/messages) share - their inner CustomStreamWrapper's logging_obj and so store + - Every other anthropic_messages/aresponses/atext_completion stream + gets a closure that dispatches on the stored args shape, because the + arming site cannot tell the producers apart: native iterators store + a single ready-made logging coroutine to enqueue, while bridged + streams (LiteLLMCompletionStreamingIterator, TextCompletionStreamWrapper, + and the plain SSE generator AnthropicStreamWrapper returns for bridged + /v1/messages) share their inner CustomStreamWrapper's logging_obj and so store (assembled_response, cache_hit); for those the closure only dispatches success logging, matching the route's pre-existing hook surface. @@ -4102,7 +4102,11 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete return - if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): + if route_type not in ( + "anthropic_messages", + "aresponses", + "atext_completion", + ) or not self._is_streaming_response(response): return from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 7ff5aa2466b..e57d6e78378 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -82,6 +82,22 @@ def _as_plain_dict(item: object) -> Mapping[str, object]: return cast("Mapping[str, object]", item) # cast-ok: wire rows are message/tool-call dicts +_STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _str_keyed(value: object) -> Mapping[str, object]: + return _STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else {} + + +def _request_route(request_data: Mapping[str, object]) -> str | None: + litellm_metadata: Final = request_data.get("litellm_metadata") + metadata_bucket: Final = _str_keyed( + litellm_metadata if litellm_metadata is not None else request_data.get("metadata") + ) + route: Final = metadata_bucket.get("user_api_key_request_route") + return route if isinstance(route, str) else None + + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -625,11 +641,8 @@ class GraySwanGuardrail(CustomGuardrail): from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings - litellm_metadata: Final = request_data.get("litellm_metadata") - request_route: Final = ( - litellm_metadata.get("user_api_key_request_route") if isinstance(litellm_metadata, Mapping) else None - ) - route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None + request_route: Final = _request_route(request_data) + route_call_types: Final = get_call_types_for_route(request_route) if request_route is not None else None call_type: Final = ( (route_call_types[0].value if route_call_types else None) or (logging_obj.call_type if logging_obj is not None else None) @@ -698,13 +711,12 @@ class GraySwanGuardrail(CustomGuardrail): inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj) - litellm_metadata: Final = request_data.get("litellm_metadata") - cleaned_litellm_metadata: Final = dict(litellm_metadata) if isinstance(litellm_metadata, dict) else {} - if inbound_headers: - existing_headers: Final = cleaned_litellm_metadata.get("headers") - cleaned_litellm_metadata["headers"] = ( - {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers - ) + litellm_metadata: Final = _str_keyed(request_data.get("litellm_metadata")) + cleaned_litellm_metadata: Final = ( + {**litellm_metadata, "headers": {**_str_keyed(litellm_metadata.get("headers")), **inbound_headers}} + if inbound_headers + else litellm_metadata + ) if cleaned_litellm_metadata: sanitized: Final = _DECODED_JSON.validate_python( safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f58c8705815..71b8976bebf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -12,11 +12,11 @@ import json import os import re import time -from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime from re import Pattern -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast import yaml from fastapi import HTTPException @@ -40,6 +40,7 @@ from litellm.proxy.guardrails.content_filter_data import ( ) from litellm.types.utils import ( CallTypes, + Delta, Function, GenericGuardrailAPIInputs, GuardrailStatus, @@ -47,6 +48,9 @@ from litellm.types.utils import ( Message, ModelResponse, ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, ) if TYPE_CHECKING: @@ -78,6 +82,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_StreamItemT: Final = TypeVar("_StreamItemT") _JSON_ENTRIES: Final = TypeAdapter(list[object]) GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b") SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+") @@ -139,6 +144,26 @@ class _CategoryConfigView(TypedDict): category_file: str | None +def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None: + if isinstance(choice, TextChoices): + return choice.text if isinstance(choice.text, str) else "" + delta: Final = getattr(choice, "delta", None) + if not delta: + return None + content: Final = getattr(delta, "content", None) + return content if isinstance(content, str) else "" + + +def _streamed_choices(item: object) -> Sequence[StreamingChoices | TextChoices]: + if isinstance(item, (ModelResponseStream, TextCompletionResponse)): + return item.choices + return () + + +def _is_choiceless_chunk(item: object) -> bool: + return isinstance(item, (ModelResponseStream, TextCompletionResponse)) and not item.choices + + @dataclass(frozen=True, slots=True) class _StreamedChoiceState: buffered_text: str = "" @@ -146,6 +171,52 @@ class _StreamedChoiceState: committed_detections: tuple[ContentFilterDetection, ...] = () latest_detections: tuple[ContentFilterDetection, ...] = () next_trim_len: int = 0 + finished: bool = False + + +def _end_of_stream_chunk( + last_item: object, state_by_choice: Mapping[int, _StreamedChoiceState] +) -> ModelResponseStream | TextCompletionResponse | None: + unfinished: Final = tuple(index for index, state in state_by_choice.items() if not state.finished) + if not unfinished: + return None + if isinstance(last_item, TextCompletionResponse): + return TextCompletionResponse( + id=last_item.id, + created=last_item.created, + model=last_item.model, + object="text_completion", + choices=[TextChoices(index=index, text="") for index in unfinished], + ).model_copy(update={"usage": None}) + if isinstance(last_item, ModelResponseStream): + return ModelResponseStream( + id=last_item.id, + created=last_item.created, + model=last_item.model, + choices=[StreamingChoices(index=index, delta=Delta(content="")) for index in unfinished], + ) + return None + + +async def _stream_then_end_of_stream_chunk( + response: AsyncIterable[_StreamItemT], end_of_stream_chunk: Callable[[_StreamItemT], _StreamItemT | None] +) -> AsyncGenerator[tuple[_StreamItemT, bool], None]: + last_item: _StreamItemT | None = None # rebind-ok: the end-of-stream chunk copies the last chunk's id and model + held_items: tuple[_StreamItemT, ...] = () # rebind-ok: the usage chunk must stay after the end-of-stream chunk + async for item in response: + last_item = item + if _is_choiceless_chunk(item): + held_items = (*held_items, item) + continue + for held_item in held_items: + yield held_item, False + held_items = () + yield item, False + closing_item: Final = None if last_item is None else end_of_stream_chunk(last_item) + if closing_item is not None: + yield closing_item, True + for held_item in held_items: + yield held_item, False @dataclass(frozen=True, slots=True) @@ -2104,72 +2175,75 @@ class ContentFilterGuardrail(CustomGuardrail): ) try: - async for item in response: - if isinstance(item, ModelResponseStream) and item.choices: - for choice in item.choices: - if not (hasattr(choice, "delta") and choice.delta): - continue + async for item, stream_ended in _stream_then_end_of_stream_chunk( + response, lambda last_item: _end_of_stream_chunk(last_item, state_by_choice) + ): + for choice in _streamed_choices(item): + new_content = _streamed_choice_content(choice) + if new_content is None: + continue - choice_index = getattr(choice, "index", 0) - if not isinstance(choice_index, int): - choice_index = 0 + choice_index = getattr(choice, "index", 0) + if not isinstance(choice_index, int): + choice_index = 0 - content = getattr(choice.delta, "content", None) - is_final = bool(getattr(choice, "finish_reason", None)) - new_content = content if isinstance(content, str) else "" - if not new_content and not is_final: - continue + is_final = stream_ended or bool(getattr(choice, "finish_reason", None)) + if not new_content and not is_final: + continue - previous_state = state_by_choice.get(choice_index, _StreamedChoiceState()) - buffered_text = previous_state.buffered_text + new_content - if not buffered_text: - continue + previous_state = state_by_choice.get(choice_index, _StreamedChoiceState()) + buffered_text = previous_state.buffered_text + new_content + if not buffered_text: + continue - # Add a space at the end if it's the final chunk to trigger word boundaries (\b) - text_to_scan = buffered_text + (" " if is_final else "") - choice_detections: list[ContentFilterDetection] = [] - scan_started = time.perf_counter() + # Add a space at the end if it's the final chunk to trigger word boundaries (\b) + text_to_scan = buffered_text + (" " if is_final else "") + choice_detections: list[ContentFilterDetection] = [] + scan_started = time.perf_counter() - try: - masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) - if is_final and masked_text.endswith(" "): - masked_text = masked_text[:-1] - latest_detections = tuple(choice_detections) - except HTTPException: - state_by_choice[choice_index] = replace( - previous_state, latest_detections=tuple(choice_detections) - ) - raise - except Exception as e: - verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) - masked_text = text_to_scan # Fallback to current text - latest_detections = previous_state.latest_detections - finally: - scan_seconds += time.perf_counter() - scan_started - - safe_to_yield_len = max( - previous_state.yielded_masked_text_len, - len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + try: + masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) + if is_final and masked_text.endswith(" "): + masked_text = masked_text[:-1] + latest_detections = tuple(choice_detections) + except HTTPException: + state_by_choice[choice_index] = replace( + previous_state, latest_detections=tuple(choice_detections) ) - choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] - next_state = replace( - previous_state, - buffered_text=buffered_text, - yielded_masked_text_len=safe_to_yield_len, - latest_detections=latest_detections, - ) - if is_final: - state_by_choice[choice_index] = next_state - continue + raise + except Exception as e: + verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) + masked_text = text_to_scan # Fallback to current text + latest_detections = previous_state.latest_detections + finally: + scan_seconds += time.perf_counter() - scan_started - trim_started = time.perf_counter() - state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan) - scan_seconds += time.perf_counter() - trim_started + safe_to_yield_len = max( + previous_state.yielded_masked_text_len, + len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + ) + released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] + if isinstance(choice, TextChoices): + choice.text = released_text + else: + choice.delta.content = released_text + next_state = replace( + previous_state, + buffered_text=buffered_text, + yielded_masked_text_len=safe_to_yield_len, + latest_detections=latest_detections, + ) + if is_final: + state_by_choice[choice_index] = replace(next_state, finished=True) + continue - yield item - else: - # Not a ModelResponseStream or no choices - yield as is - yield item + trim_started = time.perf_counter() + state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan) + scan_seconds += time.perf_counter() - trim_started + + if stream_ended and not any(_streamed_choice_content(choice) for choice in _streamed_choices(item)): + continue + yield item except HTTPException: status = "guardrail_intervened" raise diff --git a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py index 229a9be6d20..ea0a820a705 100644 --- a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py +++ b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py @@ -27,7 +27,12 @@ _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _PROVIDER_KEY: Final = "synthetic-provider-key" _GUARDRAIL: Final = "audit-text-rebuild-tool-permission" _TEXT_MODEL: Final = f"audit-text-instruct-{uuid.uuid4().hex}" +_GOVERNED_MODEL: Final = f"audit-text-governed-{uuid.uuid4().hex}" +_KEYWORD_FILTER: Final = "audit-text-keyword-filter" +_POLICY: Final = "audit-text-post-call-governance" _PIECES: Final = ("Hello", " streamed", " world") +_MASKED_PIECES: Final = ("A pers", "im", "mon tart") +_BLOCKED_PIECES: Final = ("A kum", "quat tart") _PROMPT_TOKENS: Final = 11 _COMPLETION_TOKENS: Final = 3 _USAGE: Final[dict[str, JsonValue]] = { @@ -97,6 +102,14 @@ def _text_choice(text: str, finish_reason: str | None) -> dict[str, JsonValue]: return {"text": text, "index": 0, "logprobs": None, "finish_reason": finish_reason} +def _pieces_for(identity: str) -> tuple[str, ...]: + if identity.startswith("cmpl-mask-"): + return _MASKED_PIECES + if identity.startswith("cmpl-block-"): + return _BLOCKED_PIECES + return _PIECES + + def _respond(request: Request) -> Reply: body: Final = JSON_OBJECT.validate_json(request.body) assert body.get("stream") is True, body @@ -106,7 +119,7 @@ def _respond(request: Request) -> Reply: return Reply( content_type="text/event-stream", chunks=( - *(_text_frame(identity, [_text_choice(piece, None)]) for piece in _PIECES), + *(_text_frame(identity, [_text_choice(piece, None)]) for piece in _pieces_for(identity)), _text_frame(identity, [_text_choice("", "length")]), _text_frame(identity, [], usage=_USAGE), b"data: [DONE]\n\n", @@ -114,20 +127,25 @@ def _respond(request: Request) -> Reply: ) +def _text_deployment(model_name: str, wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_base": f"{wire.url}/v1", + "api_key": _PROVIDER_KEY, + }, + } + + def _config(directory: Path, wire: Wire) -> Path: base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text())) config: Final = { **base, "model_list": [ *list_value(base["model_list"]), - { - "model_name": _TEXT_MODEL, - "litellm_params": { - "model": "text-completion-openai/gpt-3.5-turbo-instruct", - "api_base": f"{wire.url}/v1", - "api_key": _PROVIDER_KEY, - }, - }, + _text_deployment(_TEXT_MODEL, wire), + _text_deployment(_GOVERNED_MODEL, wire), ], "guardrails": [ { @@ -139,8 +157,30 @@ def _config(directory: Path, wire: Wire) -> Path: "default_action": "allow", "rules": [], }, - } + }, + { + "guardrail_name": _KEYWORD_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [ + {"keyword": "persimmon", "action": "MASK"}, + {"keyword": "kumquat", "action": "BLOCK"}, + ], + }, + }, ], + "policies": { + _POLICY: { + "guardrails": {"add": [_KEYWORD_FILTER]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": _KEYWORD_FILTER, "on_pass": "allow", "on_fail": "block"}], + }, + } + }, + "policy_attachments": [{"policy": _POLICY, "models": [_GOVERNED_MODEL]}], } path: Final = directory / f"text-rebuild-{uuid.uuid4().hex}.yaml" path.write_text(yaml.safe_dump(config)) @@ -163,14 +203,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: yield _Rig(wire, owned) -def _stream(rig: _Rig, identity: str) -> _Stream: +def _stream(rig: _Rig, identity: str, model: str = _TEXT_MODEL) -> _Stream: gateway: Final = rig.proxy.gateway with gateway.client.stream( "POST", "/v1/completions", headers={"Authorization": f"Bearer {gateway.key}"}, json={ - "model": _TEXT_MODEL, + "model": model, "prompt": identity, "max_tokens": 16, "stream": True, @@ -212,6 +252,37 @@ def test_guarded_text_completion_stream_reaches_the_caller_whole_with_its_usage( assert _GUARDRAIL in list_value(object_value(row["metadata"])["applied_guardrails"]), row +def test_post_call_pipeline_masks_a_streamed_text_completion(rig: _Rig) -> None: + identity: Final = f"cmpl-mask-{uuid.uuid4().hex}" + + stream: Final = _stream(rig, identity, _GOVERNED_MODEL) + + assert stream.status == 200, stream.body + assert stream.body.rstrip().endswith(b"data: [DONE]"), stream.body[-400:] + assert stream.text == "A [KEYWORD_REDACTED] tart", stream.frames + assert b"persimmon" not in stream.body, stream.body + assert stream.finish_reasons == ("length",), stream.frames + assert [(usage["prompt_tokens"], usage["completion_tokens"]) for usage in stream.usages] == [ + (_PROMPT_TOKENS, _COMPLETION_TOKENS) + ], stream.frames + (row,) = eventually( + lambda: rig.proxy.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (identity,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert _KEYWORD_FILTER in list_value(object_value(row["metadata"])["applied_guardrails"]), row + + +def test_post_call_pipeline_blocks_a_streamed_text_completion(rig: _Rig) -> None: + identity: Final = f"cmpl-block-{uuid.uuid4().hex}" + + stream: Final = _stream(rig, identity, _GOVERNED_MODEL) + + assert stream.status == 400, stream.body + assert b"Content blocked: keyword 'kumquat' detected" in stream.body, stream.body + assert b"quat tart" not in stream.body, stream.body + + def _together_chunk(choices: list[JsonValue], **extra: JsonValue) -> bytes: return _frame( { diff --git a/tests/unit/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py index 637a20571c8..e01c710dfdb 100644 --- a/tests/unit/litellm_core_utils/test_served_output_texts.py +++ b/tests/unit/litellm_core_utils/test_served_output_texts.py @@ -9,7 +9,16 @@ from litellm.litellm_core_utils.served_output_texts import ( served_output_texts, served_stream_output_texts, ) -from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, +) RAW = "Card: 4111 1111 1111 1111" MASKED = "Card: " @@ -83,6 +92,30 @@ def test_served_stream_output_texts_keeps_every_choice_index_when_chunks_carry_o assert served_stream_output_texts(chunks) == ("first choice", MASKED) +def _text_completion_chunk(text: str | None, index: int = 0) -> TextCompletionResponse: + return TextCompletionResponse(object="text_completion", choices=[TextChoices(index=index, text=text)]) + + +def test_served_stream_output_texts_joins_text_completion_chunks_per_choice(): + chunks = [ + _text_completion_chunk("Card: ", 0), + _text_completion_chunk("second", 1), + _text_completion_chunk("", 0), + _text_completion_chunk(None, 1), + TextCompletionResponse(object="text_completion", choices=[]), + ] + assert served_stream_output_texts(chunks) == (MASKED, "second") + + +def test_overlay_replaces_text_completion_choice_text_with_served_text(): + logged = {"id": "x", "object": "text_completion", "choices": [{"index": 0, "text": RAW}, {"index": 1, "text": RAW}]} + assert overlay_served_output_texts(logged, (MASKED,)) == { + "id": "x", + "object": "text_completion", + "choices": [{"index": 0, "text": MASKED}, {"index": 1, "text": RAW}], + } + + @pytest.mark.parametrize("content", (None, "", "first")) def test_null_choice_index_joins_choice_zero_without_changing_other_choices(content: str | None) -> None: chunks: Final = ( diff --git a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py index c96fbf34fe1..9beb1faf7f4 100644 --- a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py +++ b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py @@ -2,18 +2,25 @@ Unit tests for OpenAI Text Completion Guardrail Translation Handler """ -from typing import List, Optional, Tuple +import json +from typing import List, Literal, Optional, Tuple from unittest.mock import MagicMock import pytest +from fastapi import HTTPException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.completion.guardrail_translation.handler import ( OpenAITextCompletionHandler, ) -from litellm.types.utils import CallTypes, TextChoices, TextCompletionResponse +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage class MockGuardrail(CustomGuardrail): @@ -341,3 +348,215 @@ class TestPIIMaskingScenario: assert "[EMAIL_REDACTED]" in result["prompt"][1] assert "alice@example.com" not in result["prompt"][0] assert "bob@test.org" not in result["prompt"][1] + + +class _RecordingKeywordFilter(ContentFilterGuardrail): + def __init__(self): + super().__init__( + guardrail_name="keyword-filter", + blocked_words=[ + BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK), + BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK), + ], + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + self.streaming_sampling_rate = 1 + self.scanned_texts: tuple[tuple[str, ...], ...] = () + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.scanned_texts = (*self.scanned_texts, tuple(inputs.get("texts") or ())) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + +class _PassthroughReplaceGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if any("kumquat" in text for text in inputs.get("texts") or []): + self.raise_passthrough_exception(violation_message="replaced kumquat", request_data=request_data) + return inputs + + +def _text_chunk(index: int | None, text: str | None, finish_reason: str | None = None) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-stream", + created=1589478378, + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)], + ) + + +def _usage_chunk() -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-stream", + created=1589478378, + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[], + usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), + ) + + +def _delivered(chunks: list[TextCompletionResponse]) -> list[tuple[int, str | None, str | None]]: + return [(choice.index, choice.text, choice.finish_reason) for chunk in chunks for choice in chunk.choices] + + +async def _drive_flat_stream(guardrail: CustomGuardrail, chunks: list[TextCompletionResponse]) -> list[object]: + async def upstream(): + for chunk in chunks: + yield chunk + + return [ + item + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/completions"), + response=upstream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-3.5-turbo-instruct"}, + ) + ] + + +class TestStreamingOutputProcessing: + @pytest.mark.asyncio + async def test_ended_stream_delivers_each_choice_masked_text(self): + guardrail = _RecordingKeywordFilter() + chunks = [ + _text_chunk(0, "A pers"), + _text_chunk(1, "kiwi "), + _text_chunk(0, "im"), + _text_chunk(1, "pie", "stop"), + _text_chunk(0, "mon tart"), + _text_chunk(0, None, "length"), + _usage_chunk(), + ] + + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + request_data={}, + deliver_ended_stream_rewrites=True, + ) + + assert guardrail.scanned_texts == (("A persimmon tart", "kiwi pie"),) + assert _delivered(chunks) == [ + (0, "A [KEYWORD_REDACTED] tart", None), + (1, "kiwi ", None), + (0, "", None), + (1, "pie", "stop"), + (0, "", None), + (0, None, "length"), + ] + assert chunks[-1].usage.total_tokens == 8 + + @pytest.mark.asyncio + async def test_ended_stream_scanned_without_delivery_keeps_chunks_and_stores_the_assembled_response(self): + chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + request_data: dict[str, object] = {} + + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), request_data=request_data + ) + + assert _delivered(chunks) == [(0, "A pers", None), (0, "immon tart", None), (0, "", "stop")] + assembled = request_data["response"] + assert isinstance(assembled, TextCompletionResponse) + assert [(c.text, c.finish_reason) for c in assembled.choices] == [("A [KEYWORD_REDACTED] tart", "stop")] + assert assembled.usage.total_tokens == 8 + + @pytest.mark.asyncio + async def test_ended_stream_block_raises_the_guardrail_error(self): + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")] + + with pytest.raises(HTTPException) as exc_info: + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True + ) + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_ended_stream_treats_a_null_choice_index_as_choice_zero(self): + guardrail = _RecordingKeywordFilter() + chunks = [_text_chunk(None, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop")] + + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=guardrail, request_data={}, deliver_ended_stream_rewrites=True + ) + + assert guardrail.scanned_texts == (("A persimmon tart",),) + assert [choice.text for chunk in chunks for choice in chunk.choices] == ["A [KEYWORD_REDACTED] tart", "", ""] + + @pytest.mark.asyncio + async def test_ended_stream_blocks_a_keyword_split_across_a_null_and_a_zero_choice_index(self): + chunks = [_text_chunk(None, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")] + + with pytest.raises(HTTPException) as exc_info: + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True + ) + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_flat_stream_is_scanned_once_after_it_ends(self): + guardrail = _RecordingKeywordFilter() + chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(guardrail, chunks) + + assert guardrail.scanned_texts == (("A persimmon tart",),) + assert len(delivered) == len(chunks) + + @pytest.mark.asyncio + async def test_flat_stream_block_ends_with_an_in_band_error_frame(self): + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(_RecordingKeywordFilter(), chunks) + + error_frame = delivered[-1] + assert isinstance(error_frame, bytes) + assert error_frame.startswith(b'data: {"error": ') + assert b"kumquat" in error_frame + + @pytest.mark.asyncio + async def test_flat_stream_passthrough_block_appends_the_replacement_and_ends_with_content_filter(self): + guardrail = _PassthroughReplaceGuardrail( + guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(guardrail, chunks) + + frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)] + assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [ + ("replaced kumquat", None), + ("", "content_filter"), + ] + assert frames[-1]["usage"]["total_tokens"] == 8 + + @pytest.mark.asyncio + async def test_flat_stream_passthrough_block_reports_no_usage_the_upstream_never_sent(self): + guardrail = _PassthroughReplaceGuardrail( + guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")] + + delivered = await _drive_flat_stream(guardrail, [chunk.model_copy(update={"usage": None}) for chunk in chunks]) + + frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)] + assert frames[-1]["choices"][0]["finish_reason"] == "content_filter" + assert [frame for frame in frames if "usage" in frame] == [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index f7fb7e1fed0..9a14222c684 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -4,6 +4,7 @@ Tests for the Content Filter Guardrail import json import os +from collections.abc import AsyncGenerator, AsyncIterator, Callable from typing import Final from unittest.mock import MagicMock @@ -30,7 +31,17 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor ContentFilterCategoryConfig, ContentFilterDetection, ) -from litellm.types.utils import StandardLoggingGuardrailInformation +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field +from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StandardLoggingGuardrailInformation, + StreamingChoices, + TextChoices, + TextCompletionResponse, + Usage, +) class TestContentFilterGuardrail: @@ -3523,3 +3534,180 @@ async def test_image_whose_description_is_clean_passes_through_unchanged(): ) assert result == {"images": ["https://example.com/desk.png"], "texts": []} + + +def _fruit_word_filter() -> ContentFilterGuardrail: + return ContentFilterGuardrail( + guardrail_name="fruit-filter", + blocked_words=[ + BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK), + BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK), + ], + event_hook=GuardrailEventHooks.post_call, + ) + + +async def _text_completion_stream(*parts: tuple[str | None, str | None]): + for text, finish_reason in parts: + yield TextCompletionResponse( + id="cmpl-stream", + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)], + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_masks_a_word_split_across_text_completion_chunks(): + request_data: Final[dict] = {} + + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=_text_completion_stream(("A pers", None), ("immon tart", None), (None, "stop")), + request_data=request_data, + ) + ] + + assert "".join(chunk.choices[0].text or "" for chunk in delivered) == "A [KEYWORD_REDACTED] tart" + assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_status"] == "success" + + +_LONG_TAIL: Final = "the quick brown fox jumps over the lazy dog and then eats a ripe persimmon" + + +async def _text_completion_n2_stream_without_choice_1_finish() -> AsyncGenerator[TextCompletionResponse, None]: + for index, text, finish_reason in ( + (0, "A short answer", None), + (1, _LONG_TAIL[:30], None), + (1, _LONG_TAIL[30:], None), + (0, None, "stop"), + ): + yield TextCompletionResponse( + id="cmpl-stream", + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)], + ).model_copy(update={"usage": None}) + + +async def _chat_n2_stream_without_choice_1_finish() -> AsyncGenerator[ModelResponseStream, None]: + for index, content, finish_reason in ( + (0, "A short answer", None), + (1, _LONG_TAIL[:30], None), + (1, _LONG_TAIL[30:], None), + (0, None, "stop"), + ): + yield ModelResponseStream( + id="chatcmpl-stream", + model="gpt-4.1-mini", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + +def _streamed_text(chunk: ModelResponseStream | TextCompletionResponse, index: int) -> str: + return "".join(_choice_text(choice) for choice in chunk.choices if choice.index == index) + + +def _choice_text(choice: TextChoices | StreamingChoices) -> str: + content: Final = stream_item_field(choice, "text") if isinstance(choice, TextChoices) else choice.delta.content + return content if isinstance(content, str) else "" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stream", [_text_completion_n2_stream_without_choice_1_finish, _chat_n2_stream_without_choice_1_finish] +) +async def test_streaming_hook_releases_the_held_tail_of_a_choice_the_stream_never_finished( + stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]], +) -> None: + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=stream(), request_data={} + ) + ] + + assert "".join(_streamed_text(chunk, 0) for chunk in delivered) == "A short answer" + assert "".join(_streamed_text(chunk, 1) for chunk in delivered) == _LONG_TAIL.replace( + "persimmon", "[KEYWORD_REDACTED]" + ) + assert all(getattr(chunk, "usage", None) is None for chunk in delivered) + + +_STREAM_USAGE: Final = Usage(prompt_tokens=5, completion_tokens=20, total_tokens=25) + + +async def _then_usage_chunk( + stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]], + usage_chunk: ModelResponseStream | TextCompletionResponse, +) -> AsyncGenerator[ModelResponseStream | TextCompletionResponse, None]: + async for chunk in stream(): + yield chunk + yield usage_chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "usage_chunk"), + [ + ( + _text_completion_n2_stream_without_choice_1_finish, + TextCompletionResponse(id="cmpl-stream", object="text_completion", choices=[], usage=_STREAM_USAGE), + ), + (_chat_n2_stream_without_choice_1_finish, ModelResponseStream(id="chatcmpl-stream", choices=[], usage=_STREAM_USAGE)), + ], +) +async def test_streaming_hook_keeps_the_usage_chunk_last_after_releasing_a_held_tail( + stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]], + usage_chunk: ModelResponseStream | TextCompletionResponse, +) -> None: + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=_then_usage_chunk(stream, usage_chunk), request_data={} + ) + ] + + assert delivered[-1] is usage_chunk + assert "".join(_streamed_text(chunk, 1) for chunk in delivered[:-1]) == _LONG_TAIL.replace( + "persimmon", "[KEYWORD_REDACTED]" + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_relays_each_responses_api_event_as_it_arrives() -> None: + timeline: Final[list[str]] = [] + + async def responses_api_events() -> AsyncGenerator[OutputTextDeltaEvent, None]: + for delta in ("first", "second"): + timeline.append(f"upstream sent {delta}") + yield OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=delta, + ) + + async for event in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=responses_api_events(), request_data={} + ): + timeline.append(f"client got {event.delta}") + + assert timeline == ["upstream sent first", "client got first", "upstream sent second", "client got second"] + + +@pytest.mark.asyncio +async def test_streaming_hook_blocks_a_word_split_across_text_completion_chunks(): + with pytest.raises(HTTPException) as exc_info: + async for _ in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=_text_completion_stream(("A kum", None), ("quat tart", None), (None, "stop")), + request_data={}, + ): + pass + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py index 954c57b5cb6..a5ba9dab0bd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -5,6 +5,7 @@ import pytest from fastapi import HTTPException from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.llms.openai.completion.guardrail_translation.handler import OpenAITextCompletionHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.grayswan import grayswan as grayswan_module from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import ( @@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import ( GraySwanGuardrailAPIError, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import TextChoices, TextCompletionResponse @pytest.fixture @@ -884,3 +886,29 @@ async def test_pre_call_payload_unchanged() -> None: {"role": "user", "content": "second"}, ] assert "tools" not in payload + + +@pytest.mark.asyncio +async def test_post_call_never_forwards_the_proxy_key_metadata_of_the_request_metadata() -> None: + guardrail = _post_call_guardrail() + client = _CapturingClient() + guardrail.async_handler = client + + await OpenAITextCompletionHandler().process_output_response( + response=TextCompletionResponse( + id="cmpl-1", model="gpt-4.1-mini", choices=[TextChoices(index=0, text="hello", finish_reason="stop")] + ), + guardrail_to_apply=guardrail, + user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key", user_email="dev@example.com"), + request_data={ + "model": "gpt-4.1-mini", + "prompt": "say hello", + "metadata": { + "user_api_key_hash": "hashed-key", + "user_api_key_user_email": "dev@example.com", + "user_api_key_request_route": "/v1/completions", + }, + }, + ) + + assert "litellm_metadata" not in client.calls[0]["json"] diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 7e5d0f6342e..bd6190fee5f 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1537,6 +1537,27 @@ class TestArmDeferredStreamDispatch: assert recorded["cache_hit"] is False assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio + async def test_text_completion_stream_logs_after_the_guardrails_finish(self): + from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + self._processor()._arm_deferred_stream_dispatch( + response=TextCompletionStreamWrapper(completion_stream=MagicMock(), model="gpt-test"), + route_type="atext_completion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + @pytest.mark.asyncio async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self): """The router wraps iterators without _hidden_params in diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index acc1463b072..cfd4fe197be 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -41,6 +41,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) +from litellm.types.utils import TextChoices, TextCompletionResponse @pytest.fixture(autouse=True) @@ -1886,7 +1887,7 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator @pytest.mark.parametrize( "request_route", - ["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], + ["/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], ) def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response( make_user_api_key_auth, monkeypatch, caplog, request_route @@ -2036,6 +2037,120 @@ async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_every_cho assert _warnings(caplog) == [] +def _text_completion_chunk(text: str, finish_reason: str | None = None) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", model="m", choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)] + ) + + +def _rewritten_text_completion(response: Any) -> TextCompletionResponse: + return TextCompletionResponse( + id=response.id, + model=response.model, + choices=[TextChoices(index=0, text="[REWRITTEN] " + response.choices[0].text, finish_reason="stop")], + ) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_a_text_completion_stream( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=_rewritten_text_completion)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + chunks = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")] + auth = make_user_api_key_auth(request_route="/v1/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_logging.pre_call_hook( + user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True + ) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert isinstance(seen["response"], TextCompletionResponse) + assert seen["response"].choices[0].text == "hello world" + assert [(item.choices[0].text, item.choices[0].finish_reason) for item in delivered] == [ + ("[REWRITTEN] hello world", None), + ("", None), + ("", "stop"), + ] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + assert _warnings(caplog) == [] + + +class _FlatBlockGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + if any("kumquat" in text for text in inputs.get("texts") or []): + raise HTTPException(status_code=400, detail="flat guardrail blocked kumquat") + return inputs + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_a_flat_guardrail_block_on_a_text_completion_stream_under_a_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + flat: Final = _FlatBlockGuardrail( + guardrail_name="flat-block", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [flat, _live_scan_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + chunks: Final = [ + _text_completion_chunk("A kum"), + _text_completion_chunk("quat tart"), + _text_completion_chunk("", "stop"), + ] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + delivered: Final = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert seen["texts"] == ["A kumquat tart"] + assert isinstance(delivered[-1], bytes) + assert json.loads(delivered[-1].decode().removeprefix("data: "))["error"]["message"] == ( + "flat guardrail blocked kumquat" + ) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_live_pipeline_scans_what_a_disconnected_text_completion_client_received( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + monkeypatch.setattr(litellm, "callbacks", [_live_scan_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = {**_post_call_pipeline_data(step=_DETECT_ONLY_STEP, stream=True), "prompt": "hi"} + chunks: Final = [ + _text_completion_chunk("hello "), + _text_completion_chunk("world"), + _text_completion_chunk("", "stop"), + ] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + stream: Final = proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + first: Final = await stream.__anext__() + await stream.aclose() + + assert first.choices[0].text == "hello " + assert seen["texts"] == ["hello "] + assert _post_call_statuses(data) == ["success"] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none( proxy_logging, make_user_api_key_auth, monkeypatch @@ -2615,6 +2730,45 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated assert not any(item is chunk for item in delivered for chunk in chunks) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_modify_response_replaces_a_text_completion_stream( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep( + guardrail="gr-post", + on_pass="allow", + on_fail="modify_response", + modify_response_message="content policy block", + ) + ], + ) + data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks: Final = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + delivered: Final = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + frames: Final = [json.loads(item.decode().removeprefix("data: ")) for item in delivered] + assert seen["count"] == 1 + assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [ + ("content policy block", None), + ("", "content_filter"), + ] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse( proxy_logging, make_user_api_key_auth, monkeypatch