diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index a9fc2ea5a70..faa2776e3c3 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -205,17 +205,18 @@ class OpenAITextCompletionHandler(BaseTranslation): async def process_output_streaming_response( self, - responses_so_far: list[TextCompletionResponse], # mutable-ok: rewrites the caller's buffered chunks in place + 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[TextCompletionResponse]: - if not self._check_streaming_has_ended(responses_so_far): + ) -> 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(responses_so_far) + assembled: Final = self._assemble_ended_stream(chunks) pre_guardrail_texts: Final = _choice_texts(assembled) await self.process_output_response( response=assembled, @@ -230,16 +231,33 @@ class OpenAITextCompletionHandler(BaseTranslation): 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(responses_so_far, choice_index, text) + 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 = tuple(chunk for chunk in responses_so_far if isinstance(chunk, TextCompletionResponse)) + 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 for choice in streamed_choices if choice.finish_reason) + unfinished_indices: Final = tuple( + dict.fromkeys(choice.index for choice in streamed_choices if choice.index 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_stream_error_items( self, exc: "HTTPException", @@ -293,6 +311,10 @@ class OpenAITextCompletionHandler(BaseTranslation): choice.text = "" +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 "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index f32228a6204..67f126e2778 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -81,6 +81,21 @@ 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 _proxy_metadata(request_data: Mapping[str, object]) -> Mapping[str, object]: + litellm_metadata: Final = request_data.get("litellm_metadata") + if litellm_metadata is not None: + return _str_keyed(litellm_metadata) + metadata: Final = _str_keyed(request_data.get("metadata")) + return {key: value for key, value in metadata.items() if key.startswith("user_api_key_")} + + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -624,10 +639,7 @@ 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 - ) + request_route: Final = _proxy_metadata(request_data).get("user_api_key_request_route") route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None call_type: Final = ( (route_call_types[0].value if route_call_types else None) @@ -697,13 +709,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 - ) + proxy_metadata: Final = _proxy_metadata(request_data) + cleaned_litellm_metadata: Final = ( + {**proxy_metadata, "headers": {**_str_keyed(proxy_metadata.get("headers")), **inbound_headers}} + if inbound_headers + else proxy_metadata + ) if cleaned_litellm_metadata: sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) if isinstance(sanitized, dict) and sanitized: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py index 954c57b5cb6..f064f7fa5d5 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,24 @@ 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_on_a_text_completion_forwards_only_the_key_and_team_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"), + request_data={"model": "gpt-4.1-mini", "prompt": "say hello", "metadata": {"requester_note": "private"}}, + ) + + forwarded = client.calls[0]["json"]["litellm_metadata"] + assert forwarded["user_api_key_team_id"] == "team-456" + assert forwarded["user_api_key_key_alias"] == "ci-key" + assert "requester_note" not in forwarded 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 51d22fa93ce..52c3d71dd6a 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -2084,6 +2084,73 @@ async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_a_text_co 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