From 049f040ed2f913dfee3d5afd46488dbbad5b45b1 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Fri, 7 Aug 2026 11:21:55 -0700 Subject: [PATCH 1/9] fix(streaming): keep n>1 choices separate in stream_chunk_builder build_base_response emits a single hardcoded choice and every assembly step reads choices[0], while get_combined_content joins the delta content of every choice of every chunk into one string. A streamed n=2 call therefore came back as one choice holding both completions concatenated, which is text no model produced, plus whichever finish_reason arrived last. The reassembled response is what feeds cost and token accounting, what the caller gets with complete_response=True, and what goes into the response cache, so a cached streamed n=2 call later replays the concatenation. Assemble each choice index on its own by recursing over the chunks filtered to that index, then swap the results in after usage has been calculated. Usage is still computed over the full chunk list, so it is unchanged. A single-choice stream sees no behaviour change: the recursion only runs when the chunks carry more than one distinct index. --- litellm/main.py | 73 +++++++++++++++++++++++++++++++++ tests/test_litellm/test_main.py | 56 +++++++++++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a58..ee0c39976c8a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8700,6 +8700,66 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o if isinstance(computed_cost, (int, float)) and computed_cost > 0: setattr(usage, "cost", computed_cost) +def _streaming_choice_index(choice: Any) -> int | None: + index = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None) + return index if isinstance(index, int) else None + + +def _streaming_chunk_choices(chunk: Any) -> list[Any]: + choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) + return choices or [] + + +def _distinct_streaming_choice_indices(chunks: list[Any]) -> list[int]: + indices: Final = { + _streaming_choice_index(choice) for chunk in chunks for choice in _streaming_chunk_choices(chunk) + } + return sorted(index for index in indices if index is not None) + + +def _replace_chunk_choices(chunk: Any, choices: list[Any]) -> Any: + if not choices: + return chunk + if isinstance(chunk, dict): + return {**chunk, "choices": choices} + return chunk.model_copy(update={"choices": choices}) + + +def _chunks_for_streaming_choice_index(chunks: list[Any], index: int) -> list[Any]: + matched: Final = ( + (chunk, [c for c in _streaming_chunk_choices(chunk) if _streaming_choice_index(c) == index]) + for chunk in chunks + ) + return [ + _replace_chunk_choices(chunk, choices) + for chunk, choices in matched + if choices or not _streaming_chunk_choices(chunk) + ] + + +def _build_streaming_choices_per_index( + chunks: list[Any], + indices: list[int], + messages: list | None, + start_time, + end_time, +) -> list[Choices] | None: + built: Final = [ + stream_chunk_builder( + chunks=_chunks_for_streaming_choice_index(chunks, index), + messages=messages, + start_time=start_time, + end_time=end_time, + ) + for index in indices + ] + if any(not isinstance(response, ModelResponse) or not response.choices for response in built): + return None + return [ + cast(Choices, cast(ModelResponse, response).choices[0]).model_copy(update={"index": index}) + for index, response in zip(indices, built) + ] + def stream_chunk_builder( chunks: list, @@ -8732,6 +8792,13 @@ def stream_chunk_builder( ): # route to the text completion logic return stream_chunk_builder_text_completion(chunks=chunks, messages=messages) + choice_indices: Final = _distinct_streaming_choice_indices(chunks) + per_choice: Final = ( + _build_streaming_choices_per_index(chunks, choice_indices, messages, start_time, end_time) + if len(choice_indices) > 1 + else None + ) + model: Final = chunks[0]["model"] # Initialize the response dictionary response: Final = processor.build_base_response(chunks) @@ -8783,6 +8850,9 @@ def stream_chunk_builder( ) setattr(response, "usage", usage) + if per_choice is not None: + response.choices = per_choice + # Propagate provider_specific_fields from chunk hidden params when present. for chunk in reversed(chunks): if isinstance(chunk, dict): @@ -8961,6 +9031,9 @@ def stream_chunk_builder( setattr(response, "usage", usage) + if per_choice is not None: + response.choices = per_choice + # Propagate provider_specific_fields from the last chunk (contains provider # metadata like traffic_type set during streaming) for chunk in reversed(chunks): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 038df3656fea..1aa570e63d13 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1577,6 +1577,62 @@ async def test_async_mock_delay(): assert delay >= 0.01 +def _n_choice_stream_chunk(index, content, finish_reason=None, role=None): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + delta = {"content": content} + if role is not None: + delta["role"] = role + return ModelResponseStream( + id="chatcmpl-n2", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[StreamingChoices(index=index, delta=Delta(**delta), finish_reason=finish_reason)], + ) + + +def test_stream_chunk_builder_keeps_n_greater_than_one_choices_separate(): + from litellm import stream_chunk_builder + + chunks = [ + _n_choice_stream_chunk(0, "AAA", role="assistant"), + _n_choice_stream_chunk(1, "BBB", role="assistant"), + _n_choice_stream_chunk(0, " aaa"), + _n_choice_stream_chunk(1, " bbb"), + _n_choice_stream_chunk(0, None, finish_reason="stop"), + _n_choice_stream_chunk(1, None, finish_reason="length"), + ] + + response = stream_chunk_builder(chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert len(response.choices) == 2 + assert [choice.index for choice in response.choices] == [0, 1] + assert response.choices[0].message.content == "AAA aaa" + assert response.choices[1].message.content == "BBB bbb" + assert response.choices[0].finish_reason == "stop" + assert response.choices[1].finish_reason == "length" + + +def test_stream_chunk_builder_single_choice_is_unchanged(): + from litellm import stream_chunk_builder + + chunks = [ + _n_choice_stream_chunk(0, "AAA", role="assistant"), + _n_choice_stream_chunk(0, " aaa"), + _n_choice_stream_chunk(0, None, finish_reason="stop"), + ] + + response = stream_chunk_builder(chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert len(response.choices) == 1 + assert response.choices[0].index == 0 + assert response.choices[0].message.content == "AAA aaa" + assert response.choices[0].finish_reason == "stop" + + def test_stream_chunk_builder_thinking_blocks(): from litellm import stream_chunk_builder from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices From b33a1ebe8fd1b9cf03f19ed3074d31dcda9533db Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Tue, 25 Aug 2026 12:43:52 -0700 Subject: [PATCH 2/9] chore(streaming): give the per-index helpers real types The strict gate flagged the new helpers for ANN401 (bare `Any` on `choice`, `chunk` and the `_replace_chunk_choices` return) and ANN001 (`start_time` and `end_time` unannotated). Each one has a concrete type: the chunk and choice helpers already branch on `isinstance(..., dict)`, so the union of the pydantic model and a dict is what they actually take, and the two times are the `datetime.datetime | None` pair threaded into stream_chunk_builder. That also puts the already-imported StreamingChoices to use, clearing an F401. Net ruff difference on the file is 7 findings removed and none added. --- litellm/main.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index ee0c39976c8a..1a58997c1447 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8700,12 +8700,13 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o if isinstance(computed_cost, (int, float)) and computed_cost > 0: setattr(usage, "cost", computed_cost) -def _streaming_choice_index(choice: Any) -> int | None: + +def _streaming_choice_index(choice: StreamingChoices | dict[str, Any]) -> int | None: index = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None) return index if isinstance(index, int) else None -def _streaming_chunk_choices(chunk: Any) -> list[Any]: +def _streaming_chunk_choices(chunk: ModelResponseStream | dict[str, Any]) -> list[Any]: choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) return choices or [] @@ -8717,7 +8718,9 @@ def _distinct_streaming_choice_indices(chunks: list[Any]) -> list[int]: return sorted(index for index in indices if index is not None) -def _replace_chunk_choices(chunk: Any, choices: list[Any]) -> Any: +def _replace_chunk_choices( + chunk: ModelResponseStream | dict[str, Any], choices: list[Any] +) -> ModelResponseStream | dict[str, Any]: if not choices: return chunk if isinstance(chunk, dict): @@ -8741,8 +8744,8 @@ def _build_streaming_choices_per_index( chunks: list[Any], indices: list[int], messages: list | None, - start_time, - end_time, + start_time: datetime.datetime | None, + end_time: datetime.datetime | None, ) -> list[Choices] | None: built: Final = [ stream_chunk_builder( From 6679e9c3c0c99f873f60502d3384fd2b0b734058 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Tue, 25 Aug 2026 13:06:16 -0700 Subject: [PATCH 3/9] chore(streaming): use read-only collection types in the per-index helpers Giving the helpers concrete types traded ANN401 for LIT001, since `dict[...]` and `list[...]` are mutable collections in an annotation. Eleven of the fourteen are genuinely read-only and become `Mapping` / `Sequence`. `_replace_chunk_choices` narrowed with `isinstance(chunk, dict)`, which does not narrow a `Mapping`, so it now narrows on `ModelResponseStream` and the mapping branch stays a mapping. Same two expected inputs, same result for each. The remaining three have to be real lists and say why: two are forwarded to `stream_chunk_builder`, whose `chunks` and `messages` parameters are lists, and the third is assigned to `ModelResponse.choices`. `tests/test_litellm/test_main.py` is 98 passed before and 98 passed after. --- litellm/main.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 1a58997c1447..4bb09e5254ba 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8701,17 +8701,17 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o setattr(usage, "cost", computed_cost) -def _streaming_choice_index(choice: StreamingChoices | dict[str, Any]) -> int | None: +def _streaming_choice_index(choice: StreamingChoices | Mapping[str, Any]) -> int | None: index = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None) return index if isinstance(index, int) else None -def _streaming_chunk_choices(chunk: ModelResponseStream | dict[str, Any]) -> list[Any]: +def _streaming_chunk_choices(chunk: ModelResponseStream | Mapping[str, Any]) -> Sequence[Any]: choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) return choices or [] -def _distinct_streaming_choice_indices(chunks: list[Any]) -> list[int]: +def _distinct_streaming_choice_indices(chunks: Sequence[Any]) -> Sequence[int]: indices: Final = { _streaming_choice_index(choice) for chunk in chunks for choice in _streaming_chunk_choices(chunk) } @@ -8719,16 +8719,19 @@ def _distinct_streaming_choice_indices(chunks: list[Any]) -> list[int]: def _replace_chunk_choices( - chunk: ModelResponseStream | dict[str, Any], choices: list[Any] -) -> ModelResponseStream | dict[str, Any]: + chunk: ModelResponseStream | Mapping[str, Any], choices: Sequence[Any] +) -> ModelResponseStream | Mapping[str, Any]: if not choices: return chunk - if isinstance(chunk, dict): - return {**chunk, "choices": choices} - return chunk.model_copy(update={"choices": choices}) + # Narrow on the model rather than on dict, so the mapping branch stays a Mapping. + if isinstance(chunk, ModelResponseStream): + return chunk.model_copy(update={"choices": list(choices)}) + return {**chunk, "choices": choices} -def _chunks_for_streaming_choice_index(chunks: list[Any], index: int) -> list[Any]: +def _chunks_for_streaming_choice_index( + chunks: Sequence[Any], index: int +) -> list[Any]: # mutable-ok: fed straight to stream_chunk_builder, whose `chunks` parameter is a list matched: Final = ( (chunk, [c for c in _streaming_chunk_choices(chunk) if _streaming_choice_index(c) == index]) for chunk in chunks @@ -8741,12 +8744,12 @@ def _chunks_for_streaming_choice_index(chunks: list[Any], index: int) -> list[An def _build_streaming_choices_per_index( - chunks: list[Any], - indices: list[int], - messages: list | None, + chunks: Sequence[Any], + indices: Sequence[int], + messages: list | None, # mutable-ok: forwarded verbatim to stream_chunk_builder, whose `messages` is a list start_time: datetime.datetime | None, end_time: datetime.datetime | None, -) -> list[Choices] | None: +) -> list[Choices] | None: # mutable-ok: assigned to ModelResponse.choices, which is a list built: Final = [ stream_chunk_builder( chunks=_chunks_for_streaming_choice_index(chunks, index), From 8473655cead49fca8df0764d15d40d03acfb60ed Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Tue, 25 Aug 2026 13:13:34 -0700 Subject: [PATCH 4/9] chore(streaming): give the nine collection constructions their reasons LIT002 covers construction as well as annotation, so the lists, the local set and the two mapping literals in the per-index helpers each carry a `# mutable-ok:` reason now. Nothing moves at runtime: tests/test_litellm/test_main.py is 98 passed before and after. --- litellm/main.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 4bb09e5254ba..92c28bbcca2b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8708,11 +8708,11 @@ def _streaming_choice_index(choice: StreamingChoices | Mapping[str, Any]) -> int def _streaming_chunk_choices(chunk: ModelResponseStream | Mapping[str, Any]) -> Sequence[Any]: choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) - return choices or [] + return choices or [] # mutable-ok: the caller iterates it; an empty list is the natural absent value def _distinct_streaming_choice_indices(chunks: Sequence[Any]) -> Sequence[int]: - indices: Final = { + indices: Final = { # mutable-ok: a local set, discarded after sorting _streaming_choice_index(choice) for chunk in chunks for choice in _streaming_chunk_choices(chunk) } return sorted(index for index in indices if index is not None) @@ -8725,18 +8725,18 @@ def _replace_chunk_choices( return chunk # Narrow on the model rather than on dict, so the mapping branch stays a Mapping. if isinstance(chunk, ModelResponseStream): - return chunk.model_copy(update={"choices": list(choices)}) - return {**chunk, "choices": choices} + return chunk.model_copy(update={"choices": list(choices)}) # mutable-ok: model_copy stores what it is given, and ModelResponse.choices is a list + return {**chunk, "choices": choices} # mutable-ok: a new mapping so the caller's chunk is not mutated def _chunks_for_streaming_choice_index( chunks: Sequence[Any], index: int ) -> list[Any]: # mutable-ok: fed straight to stream_chunk_builder, whose `chunks` parameter is a list matched: Final = ( - (chunk, [c for c in _streaming_chunk_choices(chunk) if _streaming_choice_index(c) == index]) + (chunk, [c for c in _streaming_chunk_choices(chunk) if _streaming_choice_index(c) == index]) # mutable-ok: per-chunk filtered choices, consumed by the comprehension below for chunk in chunks ) - return [ + return [ # mutable-ok: stream_chunk_builder takes chunks as a list _replace_chunk_choices(chunk, choices) for chunk, choices in matched if choices or not _streaming_chunk_choices(chunk) @@ -8750,7 +8750,7 @@ def _build_streaming_choices_per_index( start_time: datetime.datetime | None, end_time: datetime.datetime | None, ) -> list[Choices] | None: # mutable-ok: assigned to ModelResponse.choices, which is a list - built: Final = [ + built: Final = [ # mutable-ok: one built response per choice index, zipped below stream_chunk_builder( chunks=_chunks_for_streaming_choice_index(chunks, index), messages=messages, @@ -8761,8 +8761,8 @@ def _build_streaming_choices_per_index( ] if any(not isinstance(response, ModelResponse) or not response.choices for response in built): return None - return [ - cast(Choices, cast(ModelResponse, response).choices[0]).model_copy(update={"index": index}) + return [ # mutable-ok: assigned to ModelResponse.choices, which is a list + cast(Choices, cast(ModelResponse, response).choices[0]).model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given for index, response in zip(indices, built) ] From e77568dca150933a605cb0e0004b88983fccabc8 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Thu, 27 Aug 2026 09:47:55 -0700 Subject: [PATCH 5/9] fix(streaming): narrow the built choices with isinstance instead of cast The per-index builder asserted its result with two nested cast() calls, which the repo's LIT006 budget counts as unchecked assertions. Move the narrowing into a small helper that isinstance-checks the response and its first choice and returns None when either is not what the merge needs, so the caller's fallback to the single-response path is reached by a real check rather than by faith. The two one-shot locals in the index/choices readers pick up the Final annotation LIT010 asks for. No behaviour change: the same inputs produce the same choices, and the same inputs fall back. Signed-off-by: Vineeth Sai --- litellm/main.py | 43 ++++++++++++++++++++++++++++++------------- 1 file changed, 30 insertions(+), 13 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 92c28bbcca2b..c87fa237fd00 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8702,12 +8702,12 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o def _streaming_choice_index(choice: StreamingChoices | Mapping[str, Any]) -> int | None: - index = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None) + index: Final = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None) return index if isinstance(index, int) else None def _streaming_chunk_choices(chunk: ModelResponseStream | Mapping[str, Any]) -> Sequence[Any]: - choices = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) + choices: Final = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) return choices or [] # mutable-ok: the caller iterates it; an empty list is the natural absent value @@ -8743,6 +8743,21 @@ def _chunks_for_streaming_choice_index( ] +def _renumbered_first_choice(response: Any, index: int) -> Choices | None: + """The single choice a per-index build produced, renumbered to that index. + + None when the build did not yield a usable non-streaming choice: `stream_chunk_builder` + is typed to admit a streaming response and a streaming choice, and either one leaves + nothing to merge. + """ + if not isinstance(response, ModelResponse) or not response.choices: + return None + first: Final = response.choices[0] + if not isinstance(first, Choices): + return None + return first.model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given + + def _build_streaming_choices_per_index( chunks: Sequence[Any], indices: Sequence[int], @@ -8750,21 +8765,23 @@ def _build_streaming_choices_per_index( start_time: datetime.datetime | None, end_time: datetime.datetime | None, ) -> list[Choices] | None: # mutable-ok: assigned to ModelResponse.choices, which is a list - built: Final = [ # mutable-ok: one built response per choice index, zipped below - stream_chunk_builder( - chunks=_chunks_for_streaming_choice_index(chunks, index), - messages=messages, - start_time=start_time, - end_time=end_time, + built: Final = [ # mutable-ok: one choice per index, returned as ModelResponse.choices + _renumbered_first_choice( + stream_chunk_builder( + chunks=_chunks_for_streaming_choice_index(chunks, index), + messages=messages, + start_time=start_time, + end_time=end_time, + ), + index, ) for index in indices ] - if any(not isinstance(response, ModelResponse) or not response.choices for response in built): + # One unusable build means there is nothing to merge, so the caller falls back to + # the single-response path rather than being handed a partial list. + if any(choice is None for choice in built): return None - return [ # mutable-ok: assigned to ModelResponse.choices, which is a list - cast(Choices, cast(ModelResponse, response).choices[0]).model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given - for index, response in zip(indices, built) - ] + return [choice for choice in built if choice is not None] # mutable-ok: assigned to ModelResponse.choices, which is a list def stream_chunk_builder( From edfac7e5a8e48da9e11df627fa58df2effd58476 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Thu, 27 Aug 2026 10:25:36 -0700 Subject: [PATCH 6/9] fix(streaming): annotate the narrowing helper with object rather than Any The helper isinstance-narrows its argument, so it does not need a dynamically typed annotation; object works and keeps the repo's ANN401 budget unchanged. Signed-off-by: Vineeth Sai --- litellm/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index c87fa237fd00..8f7639c4f107 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8743,7 +8743,7 @@ def _chunks_for_streaming_choice_index( ] -def _renumbered_first_choice(response: Any, index: int) -> Choices | None: +def _renumbered_first_choice(response: object, index: int) -> Choices | None: """The single choice a per-index build produced, renumbered to that index. None when the build did not yield a usable non-streaming choice: `stream_chunk_builder` From 4771bba8164b4ca9019eff82a3c9052f88792e36 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Thu, 27 Aug 2026 11:04:08 -0700 Subject: [PATCH 7/9] fix(streaming): drop the redundant choice narrowing ModelResponse.choices is declared list[Choices], so the second isinstance was an unnecessary check basedpyright counts against the budget; only the response itself can be the streaming variant. Signed-off-by: Vineeth Sai --- litellm/main.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 8f7639c4f107..5a298dbea5c5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8746,16 +8746,14 @@ def _chunks_for_streaming_choice_index( def _renumbered_first_choice(response: object, index: int) -> Choices | None: """The single choice a per-index build produced, renumbered to that index. - None when the build did not yield a usable non-streaming choice: `stream_chunk_builder` - is typed to admit a streaming response and a streaming choice, and either one leaves - nothing to merge. + None when the build did not yield a usable non-streaming response: `stream_chunk_builder` + is typed to admit a streaming one, which leaves nothing to merge. """ if not isinstance(response, ModelResponse) or not response.choices: return None - first: Final = response.choices[0] - if not isinstance(first, Choices): - return None - return first.model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given + # ModelResponse.choices is declared list[Choices], so the first entry needs no + # further narrowing; only the response itself can be the streaming variant. + return response.choices[0].model_copy(update={"index": index}) # mutable-ok: model_copy stores what it is given def _build_streaming_choices_per_index( From 6db73aecd8c89ecf89e0d549f7dbc146c54b6495 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 16 Sep 2026 13:13:59 -0700 Subject: [PATCH 8/9] fix(guardrails): check every Noma response choice --- .../guardrails/guardrail_hooks/noma/noma.py | 15 + .../guardrails/guardrail_hooks/test_noma.py | 2346 +---------------- 2 files changed, 84 insertions(+), 2277 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 7ef0a9f73f30..53998909e642 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -11,6 +11,7 @@ import warnings from collections.abc import AsyncGenerator from datetime import datetime +from types import MappingProxyType from typing import ( TYPE_CHECKING, Any, @@ -253,6 +254,20 @@ async def _process_llm_response_check( if not isinstance(response, litellm.ModelResponse): return None + if len(response.choices) > 1: + checked_contents: Final = tuple( + [ + await self._process_llm_response_check( + request_data, + response.model_copy(update=MappingProxyType({"choices": response.choices[index : index + 1]})), + user_auth, + event_type, + ) + for index in range(len(response.choices)) + ] + ) + return next((content for content in checked_contents if content is not None), None) + content = None for choice in response.choices: if isinstance(choice, litellm.Choices) and choice.message.content: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py index ee369f72f015..c916af4c5c59 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -1,2300 +1,92 @@ -import copy -import os -from typing import cast -from unittest.mock import AsyncMock, MagicMock, patch +import json +from collections.abc import AsyncIterator +from typing import Final import httpx import pytest +import pytest_asyncio +import respx -import litellm -from litellm import ModelResponse +from litellm import Choices, ModelResponse from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_hooks.noma import ( - NomaGuardrail, - NomaV2Guardrail, - initialize_guardrail, -) -import litellm.proxy.guardrails.guardrail_hooks.noma.noma as noma_legacy_module -from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage -from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Message +from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage, NomaGuardrail +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices -@pytest.fixture -def noma_guardrail(): - """Create a NomaGuardrail instance for testing""" - return NomaGuardrail( - api_key="test-api-key", - api_base="https://api.test.noma.security/", - application_id="test-app", - monitor_mode=False, - block_failures=True, - guardrail_name="test-noma-guardrail", - event_hook="pre_call", +@pytest_asyncio.fixture +async def guardrail() -> AsyncIterator[NomaGuardrail]: + instance: Final = NomaGuardrail( + api_key="test-key", + api_base="https://noma.test", + guardrail_name="test-noma", + event_hook="post_call", default_on=True, + anonymize_input=True, ) + async with httpx.AsyncClient() as client: + await instance.async_handler.client.aclose() + instance.async_handler.client = client + yield instance -@pytest.fixture -def mock_user_api_key_dict(): - """Create a mock UserAPIKeyAuth object""" - return UserAPIKeyAuth( - user_id="test-user-id", - user_email="test@example.com", - key_name="test-key", - key_alias=None, - team_id=None, - team_alias=None, - user_role=None, - api_key="test-api-key", - permissions={}, - models=[], - spend=0.0, - max_budget=None, - soft_budget=None, - tpm_limit=None, - rpm_limit=None, - metadata={}, - max_parallel_requests=None, - allowed_cache_controls=[], - model_spend={}, - model_max_budget={}, - ) - - -@pytest.fixture -def mock_request_data(): - """Create mock request data""" - return { - "messages": [ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": "Hello, how are you?"}, - ], - "litellm_call_id": "test-call-id", - "metadata": {"requester_ip_address": "192.168.1.1"}, - } - - -class TestNomaGuardrailConfiguration: - """Test configuration and initialization of Noma guardrail""" - - def test_legacy_guardrail_emits_deprecation_warning(self, monkeypatch): - monkeypatch.setattr( - noma_legacy_module, "_LEGACY_NOMA_DEPRECATION_WARNED", False - ) - with pytest.warns(DeprecationWarning, match="deprecated"): - NomaGuardrail(api_key="test-api-key") - - def test_init_with_config(self): - """Test initializing Noma guardrail via init_guardrails_v2""" - with patch.dict( - os.environ, - { - "NOMA_API_KEY": "test-api-key", - "NOMA_API_BASE": "https://api.test.noma.security/", - }, - ): - init_guardrails_v2( - all_guardrails=[ - { - "guardrail_name": "noma-pre-guard", - "litellm_params": { - "guardrail": "noma", - "mode": "pre_call", - "application_id": "test-app", - "monitor_mode": False, - "block_failures": True, - }, - } - ], - config_file_path="", - ) - - def test_init_with_env_vars(self): - """Test initialization with environment variables""" - with patch.dict( - os.environ, - { - "NOMA_API_KEY": "env-api-key", - "NOMA_API_BASE": "https://env.api.noma.security/", - "NOMA_APPLICATION_ID": "env-app-id", - "NOMA_MONITOR_MODE": "true", - "NOMA_BLOCK_FAILURES": "false", - }, - ): - guardrail = NomaGuardrail() - assert guardrail.api_key == "env-api-key" - assert guardrail.api_base == "https://env.api.noma.security/" - assert guardrail.application_id == "env-app-id" - assert guardrail.monitor_mode is True - assert guardrail.block_failures is False - - def test_init_with_params_override_env(self): - """Test that constructor params override environment variables""" - with patch.dict( - os.environ, - { - "NOMA_API_KEY": "env-api-key", - "NOMA_MONITOR_MODE": "true", - }, - ): - guardrail = NomaGuardrail( - api_key="param-api-key", - monitor_mode=False, - ) - assert guardrail.api_key == "param-api-key" - assert guardrail.monitor_mode is False - - def test_initialize_guardrail_function(self): - """Test the initialize_guardrail function""" - from litellm.types.guardrails import Guardrail, LitellmParams - - litellm_params = LitellmParams( - guardrail="noma", - mode="pre_call", - api_key="test-key", - api_base="https://test.api/", - application_id="test-app", - monitor_mode=True, - block_failures=False, - ) - - guardrail = Guardrail( - guardrail_name="test-guardrail", - litellm_params=litellm_params, - ) - - with patch("litellm.logging_callback_manager.add_litellm_callback") as mock_add: - result = initialize_guardrail(litellm_params, guardrail) - - assert isinstance(result, NomaGuardrail) - assert result.api_key == "test-key" - assert result.api_base == "https://test.api/" - assert result.application_id == "test-app" - assert result.monitor_mode is True - assert result.block_failures is False - mock_add.assert_called_once_with(result) - - def test_initialize_guardrail_use_v2_routes_to_noma_v2(self): - """Test migration routing: guardrail=noma + use_v2=True initializes NomaV2Guardrail.""" - from litellm.types.guardrails import Guardrail, LitellmParams - - litellm_params = LitellmParams( - guardrail="noma", - mode="pre_call", - use_v2=True, - api_key="test-key", - api_base="https://test.api/", - application_id="test-app", - ) - - guardrail = Guardrail( - guardrail_name="test-guardrail", - litellm_params=litellm_params, - ) - - with patch("litellm.logging_callback_manager.add_litellm_callback") as mock_add: - result = initialize_guardrail(litellm_params, guardrail) - - assert isinstance(result, NomaV2Guardrail) - assert result.api_key == "test-key" - assert result.api_base == "https://test.api" - assert result.application_id == "test-app" - mock_add.assert_called_once_with(result) - - -class TestNomaApplicationIdResolution: - """Tests for determining which applicationId is sent to Noma.""" - - @staticmethod - def _clone_user_auth(user_auth: UserAPIKeyAuth) -> UserAPIKeyAuth: - return UserAPIKeyAuth(**user_auth.model_dump()) - - @staticmethod - def _clone_request_data(request_data: dict) -> dict: - return copy.deepcopy(request_data) - - @staticmethod - async def _get_application_id( - guardrail: NomaGuardrail, - user_auth: UserAPIKeyAuth, - request_data: dict, - extra_data: dict, - ) -> str: - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": False, - "scanResult": [], - } - mock_response.raise_for_status = MagicMock() - mock_post = AsyncMock(return_value=mock_response) - - with patch.object(guardrail.async_handler, "post", mock_post): - await guardrail._call_noma_api( - payload={"input": []}, - llm_request_id=None, - request_data=request_data, - user_auth=user_auth, - extra_data=extra_data, - ) - - sent_payload = mock_post.call_args.kwargs["json"] - return sent_payload["x-noma-context"]["applicationId"] - - @pytest.mark.asyncio - async def test_application_id_prefers_extra_body( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - request_data = self._clone_request_data(mock_request_data) - request_data.setdefault("metadata", {}).setdefault("headers", {})[ - "x-noma-application-id" - ] = "header-app" - user_auth = self._clone_user_auth(mock_user_api_key_dict) - user_auth.key_alias = "alias-app" - - application_id = await self._get_application_id( - guardrail=noma_guardrail, - user_auth=user_auth, - request_data=request_data, - extra_data={"application_id": "dynamic-app"}, - ) - - assert application_id == "dynamic-app" - - @pytest.mark.asyncio - async def test_application_id_prefers_header_over_alias( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - request_data = self._clone_request_data(mock_request_data) - request_data.setdefault("metadata", {}).setdefault("headers", {})[ - "x-noma-application-id" - ] = "header-app" - user_auth = self._clone_user_auth(mock_user_api_key_dict) - user_auth.key_alias = "alias-app" - original_app_id = noma_guardrail.application_id - noma_guardrail.application_id = None - - try: - application_id = await self._get_application_id( - guardrail=noma_guardrail, - user_auth=user_auth, - request_data=request_data, - extra_data={}, - ) - finally: - noma_guardrail.application_id = original_app_id - - assert application_id == "header-app" - - @pytest.mark.asyncio - async def test_application_id_uses_guardrail_config_before_alias( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - request_data = self._clone_request_data(mock_request_data) - user_auth = self._clone_user_auth(mock_user_api_key_dict) - user_auth.key_alias = "alias-app" - original_app_id = noma_guardrail.application_id - noma_guardrail.application_id = "config-app" - - try: - application_id = await self._get_application_id( - guardrail=noma_guardrail, - user_auth=user_auth, - request_data=request_data, - extra_data={}, - ) - finally: - noma_guardrail.application_id = original_app_id - - assert application_id == "config-app" - - @pytest.mark.asyncio - async def test_application_id_falls_back_to_alias( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - request_data = self._clone_request_data(mock_request_data) - user_auth = self._clone_user_auth(mock_user_api_key_dict) - user_auth.key_alias = "alias-app" - original_app_id = noma_guardrail.application_id - noma_guardrail.application_id = None - - try: - application_id = await self._get_application_id( - guardrail=noma_guardrail, - user_auth=user_auth, - request_data=request_data, - extra_data={}, - ) - finally: - noma_guardrail.application_id = original_app_id - - assert application_id == "alias-app" - - @pytest.mark.asyncio - async def test_application_id_defaults_to_litellm( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - request_data = self._clone_request_data(mock_request_data) - user_auth = self._clone_user_auth(mock_user_api_key_dict) - user_auth.key_alias = None - original_app_id = noma_guardrail.application_id - noma_guardrail.application_id = None - - try: - application_id = await self._get_application_id( - guardrail=noma_guardrail, - user_auth=user_auth, - request_data=request_data, - extra_data={}, - ) - finally: - noma_guardrail.application_id = original_app_id - - assert application_id == "litellm" - - -class TestNomaBlockedMessage: - """Test the NomaBlockedMessage exception class""" - - def test_blocked_message_basic(self): - """Test basic blocked message creation""" - response = { - "aggregatedScanResult": True, - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": { - "result": True, - "probability": 0.9, - "status": "SUCCESS", - }, - "code": { - "result": False, - "probability": 0.1, - "status": "SUCCESS", - }, - }, - } - ], - } - - exception = NomaBlockedMessage(response) - assert exception.status_code == 400 - assert exception.detail["error"] == "Request blocked by Noma guardrail" - - def test_blocked_message_with_data_detection(self): - """Test blocked message with data detection""" - response = { - "aggregatedScanResult": True, - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "sensitiveData": { - "PII": { - "result": True, - "probability": 0.8, - "status": "SUCCESS", - }, - "PCI": { - "result": False, - "probability": 0, - "status": "SUCCESS", - }, - }, - }, - } - ], - } - - exception = NomaBlockedMessage(response) - assert exception.detail["error"] == "Request blocked by Noma guardrail" - - def test_blocked_message_with_topics(self): - """Test blocked message with topic guardrails""" - response = { - "aggregatedScanResult": True, - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "customLlm": { - "topic1": { - "result": True, - "probability": 0.95, - "status": "SUCCESS", - }, - "topic2": { - "result": False, - "probability": 0.2, - "status": "SUCCESS", - }, - }, - }, - } - ], - } - - exception = NomaBlockedMessage(response) - assert exception.detail["error"] == "Request blocked by Noma guardrail" - - -class TestNomaGuardrailHooks: - """Test the guardrail hook methods""" - - @pytest.mark.asyncio - async def test_pre_call_hook_allowed( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test pre-call hook when content is allowed""" - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "user", "type": "message", "results": {}}], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - result = await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=mock_request_data, - call_type="completion", - ) - - assert result == mock_request_data - mock_post.assert_called_once() - - # Verify API call details - call_args = mock_post.call_args - # Verify the URL endpoint - assert call_args.args[0].endswith("/ai-dr/v2/prompt/scan") - # Verify headers and JSON payload - if "headers" in call_args.kwargs: - headers = call_args.kwargs["headers"] - assert "Authorization" in headers - assert headers["Authorization"] == "Bearer test-api-key" - # Verify application ID is in JSON payload (not headers) - if "json" in call_args.kwargs: - json_payload = call_args.kwargs["json"] - assert "x-noma-context" in json_payload - assert json_payload["x-noma-context"]["applicationId"] == "test-app" - - @pytest.mark.asyncio - async def test_pre_call_hook_with_system_prompt( - self, noma_guardrail, mock_user_api_key_dict - ): - """Test pre-call hook includes system prompt in Noma API request""" - request_data = { - "messages": [ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": "Hello, how are you?"}, - ], - "litellm_call_id": "test-call-id", - "metadata": {"requester_ip_address": "192.168.1.1"}, - } - - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": False, # False means safe - "scanResult": [ - {"role": "system", "type": "message", "results": {}}, - {"role": "user", "type": "message", "results": {}}, - ], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - result = await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - assert result == request_data - mock_post.assert_called_once() - - # Verify the payload includes both system and user messages - call_args = mock_post.call_args - json_payload = call_args.kwargs["json"] - assert "input" in json_payload - messages = json_payload["input"] - - # Should have 2 messages: system and user - assert len(messages) == 2 - - # First message should be system - assert messages[0]["type"] == "message" - assert messages[0]["role"] == "system" - assert messages[0]["content"][0]["type"] == "input_text" - assert messages[0]["content"][0]["text"] == "You are a helpful assistant" - - # Second message should be user - assert messages[1]["type"] == "message" - assert messages[1]["role"] == "user" - assert messages[1]["content"][0]["type"] == "input_text" - assert messages[1]["content"][0]["text"] == "Hello, how are you?" - - @pytest.mark.asyncio - async def test_pre_call_hook_with_multiple_system_prompts( - self, noma_guardrail, mock_user_api_key_dict - ): - """Test pre-call hook combines multiple system prompts into single message""" - request_data = { - "messages": [ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "system", "content": "You should be polite and respectful"}, - {"role": "user", "content": "Hello, how are you?"}, - ], - "litellm_call_id": "test-call-id", - } - - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": False, - "scanResult": [ - {"role": "system", "type": "message", "results": {}}, - {"role": "user", "type": "message", "results": {}}, - ], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - result = await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - assert result == request_data - mock_post.assert_called_once() - - # Verify the payload combines system prompts into single message - call_args = mock_post.call_args - json_payload = call_args.kwargs["json"] - messages = json_payload["input"] - - # Should have 2 messages: 1 combined system and 1 user - assert len(messages) == 2 - - # First message should be system with combined content - assert messages[0]["role"] == "system" - assert messages[0]["content"][0]["type"] == "input_text" - assert ( - messages[0]["content"][0]["text"] - == "You are a helpful assistant You should be polite and respectful" - ) - - # Second message should be user - assert messages[1]["role"] == "user" - assert messages[1]["content"][0]["type"] == "input_text" - assert messages[1]["content"][0]["text"] == "Hello, how are you?" - - @pytest.mark.asyncio - async def test_pre_call_hook_blocked( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test pre-call hook when content is blocked""" - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": { - "result": True, - "probability": 0.9, - "status": "SUCCESS", - } - }, - } - ], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ): - with pytest.raises(NomaBlockedMessage) as exc_info: - await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=mock_request_data, - call_type="completion", - ) - - assert exc_info.value.status_code == 400 - - @pytest.mark.asyncio - async def test_pre_call_hook_monitor_mode( - self, mock_user_api_key_dict, mock_request_data - ): - """Test pre-call hook in monitor mode (logs but doesn't block)""" - guardrail = NomaGuardrail( - api_key="test-key", - monitor_mode=True, - guardrail_name="test-guardrail", - event_hook="pre_call", - default_on=True, - ) - - with patch.object( - guardrail, "_create_background_noma_check" - ) as mock_create_background: - # Should return immediately without waiting for API call - result = await guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=mock_request_data, - call_type="completion", - ) - - assert result == mock_request_data - # Verify background task was created - mock_create_background.assert_called_once() - - @pytest.mark.asyncio - async def test_pre_call_hook_monitor_mode_background_task_failure( - self, mock_user_api_key_dict, mock_request_data - ): - """Test pre-call hook in monitor mode when background task creation fails""" - guardrail = NomaGuardrail( - api_key="test-key", - monitor_mode=True, - guardrail_name="test-guardrail", - event_hook="pre_call", - default_on=True, - ) - - with patch.object( - guardrail, - "_create_background_noma_check", - side_effect=Exception("Task creation failed"), - ): - # Should still return successfully even if background task creation fails - result = await guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=mock_request_data, - call_type="completion", - ) - - assert result == mock_request_data - - @pytest.mark.asyncio - async def test_post_call_success_hook( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test post-call success hook""" - # Create a mock ModelResponse - response = ModelResponse( - id="test-response-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="I'm doing well, thank you!", role="assistant" - ), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) - - mock_api_response = MagicMock() - mock_api_response.json.return_value = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "assistant", "type": "message", "results": {}}], - } - mock_api_response.raise_for_status = MagicMock() - - # Update guardrail to use post_call event hook - noma_guardrail.event_hook = "post_call" - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_api_response - ) as mock_post: - result = await noma_guardrail.async_post_call_success_hook( - data=mock_request_data, - user_api_key_dict=mock_user_api_key_dict, - response=response, - ) - - assert result == response - mock_post.assert_called_once() - - @pytest.mark.asyncio - async def test_moderation_hook( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test moderation hook (during_call)""" - # Update guardrail to use during_call event hook - noma_guardrail.event_hook = "during_call" - - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "user", "type": "message", "results": {}}], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ): - result = await noma_guardrail.async_moderation_hook( - data=mock_request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - assert result == mock_request_data - - @pytest.mark.asyncio - async def test_api_failure_handling( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - with patch.object( - noma_guardrail.async_handler, - "post", - side_effect=httpx.HTTPStatusError( - "API Error", request=MagicMock(), response=MagicMock(status_code=500) - ), - ): - with pytest.raises(httpx.HTTPStatusError): - await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=mock_request_data, - call_type="completion", - ) - - @pytest.mark.asyncio - async def test_api_failure_no_block( - self, mock_user_api_key_dict, mock_request_data - ): - guardrail = NomaGuardrail( - api_key="test-key", - block_failures=False, - guardrail_name="test-guardrail", - event_hook="pre_call", - default_on=True, - ) - - with patch.object( - guardrail.async_handler, - "post", - side_effect=httpx.HTTPStatusError( - "API Error", request=MagicMock(), response=MagicMock(status_code=500) - ), - ): - result = await guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=mock_request_data, - call_type="completion", - ) - - assert result == mock_request_data - - -class TestBackgroundProcessing: - """Test the new background processing functionality""" - - @pytest.fixture - def monitor_mode_guardrail(self): - """Create a guardrail with monitor mode enabled""" - return NomaGuardrail( - api_key="test-api-key", - api_base="https://api.test.noma.security/", - application_id="test-app", - monitor_mode=True, # Enable monitor mode - block_failures=True, - guardrail_name="test-noma-guardrail", - event_hook="pre_call", - default_on=True, - ) - - @pytest.mark.asyncio - async def test_process_user_message_check_monitor_mode( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test shared helper method in monitor mode""" - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": {"result": True, "status": "SUCCESS"} - }, - } - ], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - monitor_mode_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - with patch.object( - monitor_mode_guardrail, "_handle_verdict_background" - ) as mock_handle_verdict: - result = await monitor_mode_guardrail._process_user_message_check( - mock_request_data, mock_user_api_key_dict - ) - - assert result is not None - mock_post.assert_called_once() - mock_handle_verdict.assert_called_once() - - @pytest.mark.asyncio - async def test_process_user_message_check_non_monitor_mode( - self, noma_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test shared helper method in non-monitor mode""" - mock_response = MagicMock() - mock_response.json.return_value = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "user", "type": "message", "results": {}}], - } - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - with patch.object(noma_guardrail, "_check_verdict") as mock_check_verdict: - result = await noma_guardrail._process_user_message_check( - mock_request_data, mock_user_api_key_dict - ) - - assert result is not None - mock_post.assert_called_once() - mock_check_verdict.assert_called_once() - - @pytest.mark.asyncio - async def test_process_llm_response_check_monitor_mode( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test LLM response processing in monitor mode""" - from litellm.types.utils import Choices, Message - - response = ModelResponse( - id="test-response-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="I'm doing well, thank you!", role="assistant" - ), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) - - mock_api_response = MagicMock() - mock_api_response.json.return_value = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "assistant", "type": "message", "results": {}}], - } - mock_api_response.raise_for_status = MagicMock() - - with patch.object( - monitor_mode_guardrail.async_handler, "post", return_value=mock_api_response - ) as mock_post: - with patch.object( - monitor_mode_guardrail, "_handle_verdict_background" - ) as mock_handle_verdict: - result = await monitor_mode_guardrail._process_llm_response_check( - mock_request_data, response, mock_user_api_key_dict - ) - - assert result == "I'm doing well, thank you!" - mock_post.assert_called_once() - mock_handle_verdict.assert_called_once() - - @pytest.mark.asyncio - async def test_check_user_message_background( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test background user message check method""" - with patch.object( - monitor_mode_guardrail, "_process_user_message_check" - ) as mock_process: - await monitor_mode_guardrail._check_user_message_background( - mock_request_data, mock_user_api_key_dict - ) - - mock_process.assert_called_once_with( - mock_request_data, mock_user_api_key_dict - ) - - @pytest.mark.asyncio - async def test_check_user_message_background_exception_handling( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test background user message check handles exceptions gracefully""" - with patch.object( - monitor_mode_guardrail, - "_process_user_message_check", - side_effect=Exception("API failed"), - ): - # Should not raise exception, just log error - await monitor_mode_guardrail._check_user_message_background( - mock_request_data, mock_user_api_key_dict - ) - - @pytest.mark.asyncio - async def test_check_llm_response_background( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test background LLM response check method""" - from litellm.types.utils import Choices, Message - - response = ModelResponse( - id="test-response-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message(content="Test response", role="assistant"), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) - - with patch.object( - monitor_mode_guardrail, "_process_llm_response_check" - ) as mock_process: - await monitor_mode_guardrail._check_llm_response_background( - mock_request_data, response, mock_user_api_key_dict - ) - - mock_process.assert_called_once_with( - mock_request_data, response, mock_user_api_key_dict - ) - - @pytest.mark.asyncio - async def test_handle_verdict_background_blocked(self, monitor_mode_guardrail): - """Test background verdict handling for blocked content""" - response_json = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": {"result": True, "status": "SUCCESS"} - }, - } - ], - } - - with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning: - await monitor_mode_guardrail._handle_verdict_background( - "user", "test message", response_json - ) - - mock_warning.assert_called_once() - assert "blocked user message" in mock_warning.call_args[0][0] - - @pytest.mark.asyncio - async def test_handle_verdict_background_allowed(self, monitor_mode_guardrail): - """Test background verdict handling for allowed content""" - response_json = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "assistant", "type": "message", "results": {}}], - } - - with patch("litellm._logging.verbose_proxy_logger.info") as mock_info: - await monitor_mode_guardrail._handle_verdict_background( - "assistant", "test response", response_json - ) - - mock_info.assert_called_once() - assert "allowed assistant message" in mock_info.call_args[0][0] - - @pytest.mark.asyncio - async def test_create_background_noma_check(self, monitor_mode_guardrail): - """Test background task creation""" - - async def dummy_coroutine(): - return "completed" - - with patch("asyncio.create_task") as mock_create_task: - monitor_mode_guardrail._create_background_noma_check(dummy_coroutine()) - mock_create_task.assert_called_once() - - @pytest.mark.asyncio - async def test_create_background_noma_check_exception(self, monitor_mode_guardrail): - """Test background task creation with exception handling""" - - async def dummy_coroutine(): - return "completed" - - with patch( - "asyncio.create_task", side_effect=Exception("Task creation failed") - ): - # Should not raise exception, just log error - monitor_mode_guardrail._create_background_noma_check(dummy_coroutine()) - - @pytest.mark.asyncio - async def test_moderation_hook_monitor_mode( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test moderation hook in monitor mode""" - # Update event hook to during_call - monitor_mode_guardrail.event_hook = "during_call" - - with patch.object( - monitor_mode_guardrail, "_create_background_noma_check" - ) as mock_create_background: - result = await monitor_mode_guardrail.async_moderation_hook( - data=mock_request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - assert result == mock_request_data - mock_create_background.assert_called_once() - - @pytest.mark.asyncio - async def test_post_call_success_hook_monitor_mode( - self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data - ): - """Test post-call success hook in monitor mode""" - from litellm.types.utils import Choices, Message - - # Update event hook to post_call - monitor_mode_guardrail.event_hook = "post_call" - - response = ModelResponse( - id="test-response-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message(content="Test response", role="assistant"), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) +def classify(request: httpx.Request) -> httpx.Response: + text: Final[str] = json.loads(request.content)["input"][0]["content"][0]["text"] + return httpx.Response(200, json={"aggregatedScanResult": text == "blocked", "scanResult": []}) - with patch.object( - monitor_mode_guardrail, "_create_background_noma_check" - ) as mock_create_background: - result = await monitor_mode_guardrail.async_post_call_success_hook( - data=mock_request_data, - user_api_key_dict=mock_user_api_key_dict, - response=response, - ) - - assert result == response - mock_create_background.assert_called_once() - - -class TestNomaImageProcessing: - """Test image processing functionality for multimodal content""" - - @pytest.fixture - def noma_guardrail(self): - """Create a NomaGuardrail instance for testing""" - return NomaGuardrail( - api_key="test-api-key", - api_base="https://api.test.noma.security/", - application_id="test-app", - monitor_mode=False, - block_failures=True, - guardrail_name="test-noma-guardrail", - event_hook="pre_call", - default_on=True, - ) - - @pytest.fixture - def mock_user_api_key_dict(self): - """Create a mock UserAPIKeyAuth object""" - return UserAPIKeyAuth( - user_id="test-user-id", - user_email="test@example.com", - key_name="test-key", - api_key="test-api-key", - permissions={}, - models=[], - spend=0.0, - metadata={}, - ) - - def test_extract_user_message_with_image_url(self): - """User message with only image_url becomes a single input_image content item.""" - from litellm.completion_extras.litellm_responses_transformation.transformation import ( - LiteLLMResponsesTransformationHandler, - ) - - handler = LiteLLMResponsesTransformationHandler() - messages: list[AllMessageValues] = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": "https://example.com/image.jpg"}, - } - ], - } - ] - - input_items, _ = handler.convert_chat_completion_messages_to_responses_api( - messages - ) - - assert len(input_items) == 1 - message = input_items[0]["content"] - assert message is not None - assert len(message) == 1 - assert message[0]["type"] == "input_image" - assert message[0]["image_url"] == "https://example.com/image.jpg" - - def test_extract_user_message_with_mixed_content(self): - """User message with text + image becomes input_text then input_image in content list.""" - from litellm.completion_extras.litellm_responses_transformation.transformation import ( - LiteLLMResponsesTransformationHandler, - ) - - handler = LiteLLMResponsesTransformationHandler() - messages: list[AllMessageValues] = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "What's in this image?", - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/image.jpg"}, - }, - ], - } - ] - - input_items, _ = handler.convert_chat_completion_messages_to_responses_api( - messages - ) - - # Match the original assertions: `message` is the content list - assert len(input_items) == 1 - message = input_items[0]["content"] - assert message is not None - assert len(message) == 2 - # First item should be text - assert message[0]["type"] == "input_text" - assert message[0]["text"] == "What's in this image?" - # Second item should be image - assert message[1]["type"] == "input_image" - assert message[1]["image_url"] == "https://example.com/image.jpg" - - def test_extract_user_message_with_multiple_images(self): - """User message with multiple images becomes multiple input_image items.""" - from litellm.completion_extras.litellm_responses_transformation.transformation import ( - LiteLLMResponsesTransformationHandler, - ) - - handler = LiteLLMResponsesTransformationHandler() - messages: list[AllMessageValues] = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Compare these images", - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/image1.jpg"}, - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/image2.jpg"}, - }, - ], - } +@pytest.mark.asyncio +async def test_later_choice_is_checked(guardrail: NomaGuardrail, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://noma.test/ai-dr/v2/prompt/scan").mock(side_effect=classify) + response: Final = ModelResponse( + choices=[ + Choices(index=0, message={"role": "assistant", "content": "allowed"}), + Choices(index=1, message={"role": "assistant", "content": "blocked"}), ] - - input_items, _ = handler.convert_chat_completion_messages_to_responses_api( - messages - ) - - # Match the original assertions - assert len(input_items) == 1 - message = input_items[0]["content"] - assert message is not None - assert len(message) == 3 - assert message[0]["type"] == "input_text" - assert message[1]["type"] == "input_image" - assert message[1]["image_url"] == "https://example.com/image1.jpg" - assert message[2]["type"] == "input_image" - assert message[2]["image_url"] == "https://example.com/image2.jpg" - - @pytest.mark.asyncio - async def test_pre_call_hook_with_image_content( - self, noma_guardrail, mock_user_api_key_dict - ): - """Test pre-call hook with image content""" - request_data = { - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": "https://example.com/test-image.jpg"}, - } - ], - } - ], - "litellm_call_id": "test-call-id", - "metadata": {"requester_ip_address": "192.168.1.1"}, - } - - # Mock Noma API response for image content - noma_response = { - "aggregatedScanResult": False, # False means safe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": { - "result": False, - "probability": 0.1, - "status": "SUCCESS", - } - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - result = await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - assert result == request_data - mock_post.assert_called_once() - - # Verify the API call payload includes image - call_args = mock_post.call_args - payload = call_args[1]["json"] - assert "input" in payload - assert len(payload["input"]) > 0 - assert payload["input"][0]["role"] == "user" - assert "content" in payload["input"][0] - - @pytest.mark.asyncio - async def test_pre_call_hook_with_mixed_content( - self, noma_guardrail, mock_user_api_key_dict - ): - """Test pre-call hook with mixed text and image content""" - request_data = { - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Analyze this image for harmful content", - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/test-image.jpg"}, - }, - ], - } - ], - "litellm_call_id": "test-call-id", - } - - # Mock Noma API response - noma_response = { - "aggregatedScanResult": False, - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": { - "result": False, - "probability": 0.05, - "status": "SUCCESS", - } - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: - result = await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - assert result == request_data - mock_post.assert_called_once() - - @pytest.mark.asyncio - async def test_image_content_blocked(self, noma_guardrail, mock_user_api_key_dict): - """Test that image content can be blocked by Noma""" - request_data = { - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "https://example.com/inappropriate-image.jpg" - }, - } - ], - } - ], - "litellm_call_id": "test-call-id", - } - - # Mock Noma API response indicating harmful content in image - noma_response = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "harmfulContent": { - "result": True, - "probability": 0.95, - "status": "SUCCESS", - } - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - from litellm.proxy.guardrails.guardrail_hooks.noma.noma import ( - NomaBlockedMessage, - ) - - with patch.object( - noma_guardrail.async_handler, "post", return_value=mock_response - ): - with pytest.raises(NomaBlockedMessage) as exc_info: - await noma_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - assert exc_info.value.status_code == 400 - - @pytest.mark.asyncio - async def test_image_with_base64_data(self, noma_guardrail, mock_user_api_key_dict): - """Test extracting image with base64 data URL""" - from litellm.completion_extras.litellm_responses_transformation.transformation import ( - LiteLLMResponsesTransformationHandler, - ) - - data = { - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - }, - } - ], - } - ] - } - - handler = LiteLLMResponsesTransformationHandler() - messages = cast(list[AllMessageValues], data["messages"]) - - input_items, _ = handler.convert_chat_completion_messages_to_responses_api( - messages - ) - - assert len(input_items) == 1 - message = input_items[0]["content"] - assert message is not None - assert len(message) == 1 - assert message[0]["type"] == "input_image" - assert message[0]["image_url"].startswith("data:image/jpeg;base64,") - - -class TestIntegration: - @pytest.mark.asyncio - async def test_full_guardrail_flow(self): - """Test full guardrail flow with multiple hooks""" - with patch.dict( - os.environ, - { - "NOMA_API_KEY": "test-api-key", - "NOMA_API_BASE": "https://api.test.noma.security/", - }, - ): - init_guardrails_v2( - all_guardrails=[ - { - "guardrail_name": "noma-pre-guard", - "litellm_params": { - "guardrail": "noma", - "mode": "pre_call", - "application_id": "test-app", - }, - }, - { - "guardrail_name": "noma-post-guard", - "litellm_params": { - "guardrail": "noma", - "mode": "post_call", - "application_id": "test-app", - }, - }, + ) + with pytest.raises(NomaBlockedMessage): + await guardrail.async_post_call_success_hook({}, UserAPIKeyAuth(), response) + assert route.call_count == 2 + + +@pytest.mark.asyncio +async def test_anonymization_stays_with_its_choice(guardrail: NomaGuardrail, respx_mock: respx.MockRouter) -> None: + def anonymize(request: httpx.Request) -> httpx.Response: + text: Final[str] = json.loads(request.content)["input"][0]["content"][0]["text"] + return httpx.Response( + 200, + json={ + "aggregatedScanResult": False, + "scanResult": [ + {"role": "assistant", "results": {"anonymizedContent": {"anonymized": text + " redacted"}}} ], - config_file_path="", - ) - - custom_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=litellm.integrations.custom_guardrail.CustomGuardrail - ) - ) - assert len(custom_loggers) >= 2 - - -class TestNomaAnonymizationConfiguration: - """Test anonymize_input configuration parameter""" - - def test_init_with_anonymize_input_env_var(self): - """Test initialization with NOMA_ANONYMIZE_INPUT environment variable""" - with patch.dict( - os.environ, - { - "NOMA_ANONYMIZE_INPUT": "true", - }, - ): - guardrail = NomaGuardrail() - assert guardrail.anonymize_input is True - - with patch.dict( - os.environ, - { - "NOMA_ANONYMIZE_INPUT": "false", - }, - ): - guardrail = NomaGuardrail() - assert guardrail.anonymize_input is False - - def test_init_with_anonymize_input_default(self): - """Test default value for anonymize_input""" - guardrail = NomaGuardrail() - assert guardrail.anonymize_input is False - - def test_init_with_anonymize_input_param_override_env(self): - """Test that constructor param overrides environment variable""" - with patch.dict( - os.environ, - { - "NOMA_ANONYMIZE_INPUT": "true", - }, - ): - guardrail = NomaGuardrail(anonymize_input=False) - assert guardrail.anonymize_input is False - - def test_initialize_guardrail_with_anonymize_input(self): - """Test the initialize_guardrail function with anonymize_input""" - from litellm.types.guardrails import Guardrail, LitellmParams - - litellm_params = LitellmParams( - guardrail="noma", - mode="pre_call", - api_key="test-key", - anonymize_input=True, - ) - - guardrail = Guardrail( - guardrail_name="test-guardrail", - litellm_params=litellm_params, - ) - - with patch("litellm.logging_callback_manager.add_litellm_callback"): - result = initialize_guardrail(litellm_params, guardrail) - assert result.anonymize_input is True - - -class TestNomaAnonymizationLogic: - """Test the anonymization logic helper methods""" - - @pytest.fixture - def anonymize_guardrail(self): - """Create a guardrail with anonymize_input enabled""" - return NomaGuardrail( - api_key="test-api-key", - anonymize_input=True, - monitor_mode=False, - block_failures=True, - ) - - def test_is_result_true(self, anonymize_guardrail): - """Test _is_result_true helper method""" - assert anonymize_guardrail._is_result_true({"result": True}) is True - assert anonymize_guardrail._is_result_true({"result": False}) is False - assert anonymize_guardrail._is_result_true({"other": True}) is False - assert anonymize_guardrail._is_result_true(None) is False - assert anonymize_guardrail._is_result_true({}) is False - assert anonymize_guardrail._is_result_true("not a dict") is False - - def test_should_only_data_detector_failed_true(self, anonymize_guardrail): - """Test _should_only_sensitive_data_failed when only sensitive data detector triggered""" - classification = { - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - "PCI": {"result": True, "status": "SUCCESS"}, - "secrets": {"result": False, "probability": 0, "status": "SUCCESS"}, - }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - "maliciousIntent": {"result": False, "status": "SUCCESS"}, - "code": {"result": False, "status": "SUCCESS"}, - } - - result = anonymize_guardrail._should_only_sensitive_data_failed(classification) - assert result is True - - def test_should_only_data_detector_failed_false_other_detectors( - self, anonymize_guardrail - ): - """Test _should_only_sensitive_data_failed when other detectors also triggered""" - classification = { - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - }, - "harmfulContent": { - "result": True, - "status": "SUCCESS", - }, # This should cause False - "maliciousIntent": {"result": False, "status": "SUCCESS"}, - } - - result = anonymize_guardrail._should_only_sensitive_data_failed(classification) - assert result is False - - def test_should_only_data_detector_failed_false_no_data_detected( - self, anonymize_guardrail - ): - """Test _should_only_sensitive_data_failed when no sensitive data detected""" - classification = { - "sensitiveData": { - "PII": {"result": False, "status": "SUCCESS"}, - "PCI": {"result": False, "status": "SUCCESS"}, }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - "maliciousIntent": {"result": False, "status": "SUCCESS"}, - } - - result = anonymize_guardrail._should_only_sensitive_data_failed(classification) - assert result is False - - def test_should_only_data_detector_failed_with_nested_detectors( - self, anonymize_guardrail - ): - """Test _should_only_sensitive_data_failed with nested detectors like topicDetector""" - classification = { - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - }, - "customLlm": { - "topic1": { - "result": True, - "status": "SUCCESS", - }, # This should cause False - }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - } - - result = anonymize_guardrail._should_only_sensitive_data_failed(classification) - assert result is False - - def test_extract_anonymized_content_user(self, anonymize_guardrail): - """Test _extract_anonymized_content for user messages""" - response_json = { - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "anonymizedContent": { - "anonymized": "My email is ******* and phone is *******" - } - }, - } - ] - } - - result = anonymize_guardrail._extract_anonymized_content(response_json, "user") - assert result == "My email is ******* and phone is *******" - - def test_extract_anonymized_content_assistant(self, anonymize_guardrail): - """Test _extract_anonymized_content for assistant messages""" - response_json = { - "scanResult": [ - { - "role": "assistant", - "type": "message", - "results": { - "anonymizedContent": { - "anonymized": "I can't help with that request." - } - }, - } - ] - } - - result = anonymize_guardrail._extract_anonymized_content( - response_json, "assistant" - ) - assert result == "I can't help with that request." - - def test_extract_anonymized_content_missing(self, anonymize_guardrail): - """Test _extract_anonymized_content when anonymized content is missing""" - response_json = { - "scanResult": [{"role": "user", "type": "message", "results": {}}] - } - - result = anonymize_guardrail._extract_anonymized_content(response_json, "user") - assert result == "" - - def test_should_anonymize_verdict_true(self, anonymize_guardrail): - """Test _should_anonymize when aggregatedScanResult is False (safe)""" - response_json = { - "aggregatedScanResult": False, # False means safe - "scanResult": [{"role": "user", "type": "message", "results": {}}], - } - - result = anonymize_guardrail._should_anonymize(response_json, "user") - assert result is True - - def test_should_anonymize_verdict_false_only_sensitive(self, anonymize_guardrail): - """Test _should_anonymize when aggregatedScanResult is True but only sensitive data detector triggered""" - response_json = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "sensitiveData": {"PCI": {"result": True, "status": "SUCCESS"}}, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - }, - } - ], - } - - result = anonymize_guardrail._should_anonymize(response_json, "user") - assert result is True - - def test_should_anonymize_verdict_false_other_detectors(self, anonymize_guardrail): - """Test _should_anonymize when aggregatedScanResult is True and other detectors triggered""" - response_json = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "sensitiveData": {"PCI": {"result": True, "status": "SUCCESS"}}, - "harmfulContent": {"result": True, "status": "SUCCESS"}, - }, - } - ], - } - - result = anonymize_guardrail._should_anonymize(response_json, "user") - assert result is False - - def test_should_anonymize_monitor_mode(self): - """Test _should_anonymize in monitor mode (should never anonymize)""" - guardrail = NomaGuardrail( - anonymize_input=True, - monitor_mode=True, - ) - - response_json = { - "aggregatedScanResult": False, - "scanResult": [{"role": "user", "type": "message", "results": {}}], - } - result = guardrail._should_anonymize(response_json, "user") - assert result is False - - def test_should_anonymize_disabled(self): - """Test _should_anonymize when anonymize_input is disabled""" - guardrail = NomaGuardrail( - anonymize_input=False, - monitor_mode=False, - ) - - response_json = { - "aggregatedScanResult": False, - "scanResult": [{"role": "user", "type": "message", "results": {}}], - } - result = guardrail._should_anonymize(response_json, "user") - assert result is False - - def test_replace_user_message_content(self, anonymize_guardrail): - """Test _replace_user_message_content""" - request_data = { - "messages": [ - {"role": "system", "content": "System prompt"}, - {"role": "user", "content": "My email is test@example.com"}, - {"role": "assistant", "content": "I can help you"}, - {"role": "user", "content": "My phone is 123-456-7890"}, - ] - } - - anonymize_guardrail._replace_user_message_content( - request_data, "My phone is *******" - ) - - # Should replace the last user message - assert request_data["messages"][-1]["content"] == "My phone is *******" - assert ( - request_data["messages"][1]["content"] == "My email is test@example.com" - ) # Unchanged - - def test_replace_llm_response_content(self, anonymize_guardrail): - """Test _replace_llm_response_content""" - response = ModelResponse( - id="test-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="Your email is test@example.com", role="assistant" - ), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) - - anonymize_guardrail._replace_llm_response_content( - response, "Your email is *******" ) - assert response.choices[0].message.content == "Your email is *******" - - -class TestNomaAnonymizationFlow: - """Test full anonymization flow with real Noma response objects""" - - @pytest.fixture - def anonymize_guardrail(self): - """Create a guardrail with anonymize_input enabled""" - return NomaGuardrail( - api_key="test-api-key", - api_base="https://api.test.noma.security/", - application_id="test-app", - anonymize_input=True, - monitor_mode=False, - block_failures=True, - guardrail_name="test-noma-guardrail", - event_hook="pre_call", - default_on=True, - ) - - @pytest.fixture - def mock_user_api_key_dict(self): - """Create a mock UserAPIKeyAuth object""" - return UserAPIKeyAuth( - user_id="test-user-id", - user_email="test@example.com", - key_name="test-key", - api_key="test-api-key", - permissions={}, - models=[], - spend=0.0, - metadata={}, - ) - - @pytest.mark.asyncio - async def test_anonymization_verdict_true_user_message( - self, anonymize_guardrail, mock_user_api_key_dict - ): - """Test anonymization when verdict=True for user message""" - request_data = { - "messages": [ - {"role": "user", "content": "My email is test@example.com"}, - ], - "litellm_call_id": "test-call-id", - "metadata": {"requester_ip_address": "192.168.1.1"}, - } - - # Mock simplified Noma API response with aggregatedScanResult=False (safe) and anonymized content - noma_response = { - "aggregatedScanResult": False, # False means safe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "anonymizedContent": {"anonymized": "My email is *******"}, - "sensitiveData": { - "PII": {"result": False, "status": "SUCCESS"}, - }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - with patch.object( - anonymize_guardrail.async_handler, "post", return_value=mock_response - ): - result = await anonymize_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - # Should return modified request with anonymized content - assert result == request_data - assert result["messages"][0]["content"] == "My email is *******" - - @pytest.mark.asyncio - async def test_anonymization_verdict_false_only_data_detected( - self, anonymize_guardrail, mock_user_api_key_dict - ): - """Test anonymization when verdict=False but only data detector triggered""" - request_data = { - "messages": [ - {"role": "user", "content": "My email is test@example.com"}, - ], - "litellm_call_id": "test-call-id", - } - - # Mock simplified Noma API response - only sensitive data detector triggered - noma_response = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "anonymizedContent": {"anonymized": "My email is *******"}, - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - "maliciousIntent": {"result": False, "status": "SUCCESS"}, - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - with patch.object( - anonymize_guardrail.async_handler, "post", return_value=mock_response - ): - result = await anonymize_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - # Should return modified request with anonymized content (not blocked) - assert result == request_data - assert result["messages"][0]["content"] == "My email is *******" - - @pytest.mark.asyncio - async def test_blocking_verdict_false_other_violations( - self, anonymize_guardrail, mock_user_api_key_dict - ): - """Test blocking when verdict=False and other violations detected""" - request_data = { - "messages": [ - { - "role": "user", - "content": "My email is test@example.com. Tell me harmful content.", - }, - ], - "litellm_call_id": "test-call-id", - } - - # Mock simplified Noma API response - both sensitive data detector and other violations - noma_response = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "anonymizedContent": { - "anonymized": "My email is *******. Tell me harmful content." - }, - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - }, - "harmfulContent": { - "result": True, - "status": "SUCCESS", - }, # This should cause blocking - "maliciousIntent": {"result": False, "status": "SUCCESS"}, - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - with patch.object( - anonymize_guardrail.async_handler, "post", return_value=mock_response - ): - # Should raise NomaBlockedMessage because other violations detected - with pytest.raises(NomaBlockedMessage) as exc_info: - await anonymize_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - assert exc_info.value.status_code == 400 - - @pytest.mark.asyncio - async def test_anonymization_llm_response( - self, anonymize_guardrail, mock_user_api_key_dict - ): - """Test anonymization of LLM response""" - request_data = { - "messages": [{"role": "user", "content": "What's your email?"}], - "litellm_call_id": "test-call-id", - } - - # Create LLM response with test data - llm_response = ModelResponse( - id="test-response-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="My email is admin@company.com", role="assistant" - ), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) - - # Mock simplified Noma API response for LLM response check - noma_response = { - "aggregatedScanResult": True, - "scanResult": [ - { - "role": "assistant", - "type": "message", - "results": { - "anonymizedContent": {"anonymized": "My email is *******"}, - "sensitiveData": { - "PCI": { - "probability": 0.8, - "result": True, - "status": "SUCCESS", - }, - }, - }, - }, - ], - } - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - # Update guardrail to use post_call event hook - anonymize_guardrail.event_hook = "post_call" - - with patch.object( - anonymize_guardrail.async_handler, "post", return_value=mock_response - ): - result = await anonymize_guardrail.async_post_call_success_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - response=llm_response, - ) - - # Should return modified response with anonymized content - assert result == llm_response - assert result.choices[0].message.content == "My email is *******" - - @pytest.mark.asyncio - async def test_no_anonymization_when_disabled(self, mock_user_api_key_dict): - """Test that no anonymization occurs when anonymize_input=False""" - guardrail = NomaGuardrail( - api_key="test-api-key", - anonymize_input=False, # Disabled - monitor_mode=False, - block_failures=True, - ) - - request_data = { - "messages": [ - {"role": "user", "content": "My email is test@example.com"}, - ], - } - - noma_response = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "anonymizedContent": {"anonymized": "My email is *******"}, - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() + route: Final = respx_mock.post("https://noma.test/ai-dr/v2/prompt/scan").mock(side_effect=anonymize) + response: Final = ModelResponse( + choices=[ + Choices(index=0, message={"role": "assistant", "content": "first"}), + Choices(index=1, message={"role": "assistant", "content": "second"}), + Choices(index=2, message={"role": "assistant", "content": None}), + ] + ) + result: Final = await guardrail.async_post_call_success_hook({}, UserAPIKeyAuth(), response) + assert [choice.message.content for choice in result.choices] == ["first redacted", "second redacted", None] + assert route.call_count == 2 - with patch.object(guardrail.async_handler, "post", return_value=mock_response): - # Should raise NomaBlockedMessage because anonymization is disabled - with pytest.raises(NomaBlockedMessage): - await guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - @pytest.mark.asyncio - async def test_no_anonymization_in_monitor_mode(self, mock_user_api_key_dict): - """Test that no anonymization occurs in monitor mode""" - guardrail = NomaGuardrail( - api_key="test-api-key", - anonymize_input=True, - monitor_mode=True, # Monitor mode - block_failures=True, - ) +@pytest.mark.asyncio +async def test_stream_checks_later_choice_before_yielding( + guardrail: NomaGuardrail, respx_mock: respx.MockRouter +) -> None: + respx_mock.post("https://noma.test/ai-dr/v2/prompt/scan").mock(side_effect=classify) - request_data = { - "messages": [ - {"role": "user", "content": "My email is test@example.com"}, - ], - } - - with patch.object( - guardrail, "_create_background_noma_check" - ) as mock_create_background: - result = await guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", + async def chunks() -> AsyncIterator[ModelResponseStream]: + for index, text in enumerate(("allowed", "blocked")): + yield ModelResponseStream( + choices=[StreamingChoices(index=index, delta=Delta(content=text), finish_reason="stop")] ) - # Should return original data unchanged - assert result == request_data - assert ( - request_data["messages"][0]["content"] == "My email is test@example.com" - ) - mock_create_background.assert_called_once() - - @pytest.mark.asyncio - async def test_anonymization_no_anonymized_content_available( - self, anonymize_guardrail, mock_user_api_key_dict - ): - """Test behavior when anonymized content is not available""" - request_data = { - "messages": [ - {"role": "user", "content": "My email is test@example.com"}, - ], - } - - noma_response = { - "aggregatedScanResult": True, # True means unsafe - "scanResult": [ - { - "role": "user", - "type": "message", - "results": { - "sensitiveData": { - "PII": {"result": True, "status": "SUCCESS"}, - }, - "harmfulContent": {"result": False, "status": "SUCCESS"}, - }, - } - ], - } - - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - with patch.object( - anonymize_guardrail.async_handler, "post", return_value=mock_response - ): - # Should raise NomaBlockedMessage because no anonymized content available - with pytest.raises(NomaBlockedMessage): - await anonymize_guardrail.async_pre_call_hook( - user_api_key_dict=mock_user_api_key_dict, - cache=MagicMock(), - data=request_data, - call_type="completion", - ) - - @pytest.mark.asyncio - async def test_anonymization_llm_response_no_anonymized_content_available( - self, anonymize_guardrail, mock_user_api_key_dict - ): - """Test behavior when LLM response has no anonymized content available""" - request_data = { - "messages": [{"role": "user", "content": "What's your email?"}], - "litellm_call_id": "test-call-id", - } - - # Create LLM response with test data - llm_response = ModelResponse( - id="test-response-id", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="My email is admin@company.com", role="assistant" - ), - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - system_fingerprint=None, - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - ) - - # Mock Noma API response with no anonymized content available - noma_response = { - "aggregatedScanResult": True, - "scanResult": [ - { - "role": "assistant", - "type": "message", - "results": { - "sensitiveData": { - "PCI": { - "probability": 0.8, - "result": True, - "status": "SUCCESS", - }, - }, - }, - }, - ], - } - mock_response = MagicMock() - mock_response.json.return_value = noma_response - mock_response.raise_for_status = MagicMock() - - # Update guardrail to use post_call event hook - anonymize_guardrail.event_hook = "post_call" - - with patch.object( - anonymize_guardrail.async_handler, "post", return_value=mock_response - ): - # Should raise NomaBlockedMessage because no anonymized content available for LLM response - with pytest.raises(NomaBlockedMessage): - await anonymize_guardrail.async_post_call_success_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - response=llm_response, - ) + stream: Final = guardrail.async_post_call_streaming_iterator_hook(UserAPIKeyAuth(), chunks(), {}) + with pytest.raises(NomaBlockedMessage): + await anext(stream) From 93c8d9e54acd0ee76c4c7f86305e3ad6072bbdd5 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 16 Sep 2026 13:26:58 -0700 Subject: [PATCH 9/9] test(guardrails): isolate the Noma test HTTP client --- .../proxy/guardrails/guardrail_hooks/test_noma.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py index c916af4c5c59..5a7c413bed00 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -8,6 +8,7 @@ import respx from litellm import Choices, ModelResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage, NomaGuardrail from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices @@ -23,9 +24,11 @@ async def guardrail() -> AsyncIterator[NomaGuardrail]: default_on=True, anonymize_input=True, ) + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() async with httpx.AsyncClient() as client: - await instance.async_handler.client.aclose() - instance.async_handler.client = client + handler.client = client + instance.async_handler = handler yield instance