From ffe5d303e5294cdcfb0db43d32e3ca385f141a41 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Fri, 18 Sep 2026 01:30:46 -0700 Subject: [PATCH 1/5] fix(llmguard): accept proxy async call types --- .../enterprise_callbacks/llm_guard.py | 18 +++- tests/local_testing/test_llm_guard.py | 93 +++++++++++++++++++ 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index d10b5a2ab09..9c8537e6820 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -8,7 +8,7 @@ ## This provides an LLM Guard Integration for content moderation on the proxy import asyncio -from typing import Optional +from typing import Final, Optional import aiohttp from fastapi import HTTPException @@ -137,15 +137,25 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return self.print_verbose("Makes LLM Guard Check") - if call_type not in [ + accepted_call_types: Final = ( "completion", + "acompletion", + "text_completion", + "atext_completion", "embeddings", + "embedding", + "aembedding", "image_generation", + "aimage_generation", "moderation", + "amoderation", "audio_transcription", - ]: + "transcription", + "atranscription", + ) + if call_type not in accepted_call_types: self.print_verbose( - f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']" + f"Call Type - {call_type}, not in accepted list - {accepted_call_types}" ) return data diff --git a/tests/local_testing/test_llm_guard.py b/tests/local_testing/test_llm_guard.py index 9e70d48dbda..ceb77386349 100644 --- a/tests/local_testing/test_llm_guard.py +++ b/tests/local_testing/test_llm_guard.py @@ -5,6 +5,7 @@ ## Unit test for presidio pii masking import sys, os, asyncio, time, random from datetime import datetime +from typing import Final, Literal import traceback from dotenv import load_dotenv @@ -19,6 +20,7 @@ from litellm import Router, mock_completion from litellm.proxy.utils import ProxyLogging, hash_token from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache +from litellm.types.utils import CallTypesLiteral ### UNIT TESTS FOR LLM GUARD ### @@ -106,6 +108,97 @@ async def test_llm_guard_sanitizes_multimodal_and_input(): assert input_result["input"] == ["email: [REDACTED]", "email: [REDACTED]"] +@pytest.mark.parametrize( + "call_type, payload_key", + ( + ("completion", "messages"), + ("acompletion", "messages"), + ("text_completion", "prompt"), + ("atext_completion", "prompt"), + ("embeddings", "input"), + ("embedding", "input"), + ("aembedding", "input"), + ("moderation", "input"), + ("amoderation", "input"), + ("image_generation", "prompt"), + ("aimage_generation", "prompt"), + ("audio_transcription", "prompt"), + ("transcription", "prompt"), + ("atranscription", "prompt"), + ), +) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_call_type_aliases( + call_type: CallTypesLiteral, + payload_key: Literal["messages", "input", "prompt"], + is_valid: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={ + "sanitized_prompt": "email: [REDACTED]", + "is_valid": is_valid, + }, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-12345")) + data: Final = { + payload_key: [{"role": "user", "content": "email: person@example.com"}] + if payload_key == "messages" + else "email: person@example.com" + } + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=call_type + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated content safety policy"} + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=call_type + ) + assert result is data + assert data[payload_key] == ( + [{"role": "user", "content": "email: [REDACTED]"}] + if payload_key == "messages" + else "email: [REDACTED]" + ) + + +@pytest.mark.parametrize( + "call_type", + ( + "responses", + "aresponses", + "anthropic_messages", + "aanthropic_messages", + "aspeech", + "aimage_edit", + "pass_through_endpoint", + ), +) +@pytest.mark.asyncio +async def test_llm_guard_skips_unsupported_call_types( + call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"is_valid": False}, + ) + data: Final = {"messages": [{"role": "user", "content": "unchanged"}]} + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data == {"messages": [{"role": "user", "content": "unchanged"}]} + + @pytest.mark.asyncio async def test_llm_guard_error_raising(): """ From 19cb6b855b606c7a86b65cbcd933b1e12a935a6a Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 07:15:51 +0000 Subject: [PATCH 2/5] test(llmguard): move call type alias tests to the mapped enterprise test file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/local_testing/test_llm_guard.py | 93 ------------------ .../enterprise_callbacks/test_llm_guard.py | 97 +++++++++++++++++++ 2 files changed, 97 insertions(+), 93 deletions(-) create mode 100644 tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py diff --git a/tests/local_testing/test_llm_guard.py b/tests/local_testing/test_llm_guard.py index ceb77386349..9e70d48dbda 100644 --- a/tests/local_testing/test_llm_guard.py +++ b/tests/local_testing/test_llm_guard.py @@ -5,7 +5,6 @@ ## Unit test for presidio pii masking import sys, os, asyncio, time, random from datetime import datetime -from typing import Final, Literal import traceback from dotenv import load_dotenv @@ -20,7 +19,6 @@ from litellm import Router, mock_completion from litellm.proxy.utils import ProxyLogging, hash_token from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache -from litellm.types.utils import CallTypesLiteral ### UNIT TESTS FOR LLM GUARD ### @@ -108,97 +106,6 @@ async def test_llm_guard_sanitizes_multimodal_and_input(): assert input_result["input"] == ["email: [REDACTED]", "email: [REDACTED]"] -@pytest.mark.parametrize( - "call_type, payload_key", - ( - ("completion", "messages"), - ("acompletion", "messages"), - ("text_completion", "prompt"), - ("atext_completion", "prompt"), - ("embeddings", "input"), - ("embedding", "input"), - ("aembedding", "input"), - ("moderation", "input"), - ("amoderation", "input"), - ("image_generation", "prompt"), - ("aimage_generation", "prompt"), - ("audio_transcription", "prompt"), - ("transcription", "prompt"), - ("atranscription", "prompt"), - ), -) -@pytest.mark.parametrize("is_valid", (True, False)) -@pytest.mark.asyncio -async def test_llm_guard_call_type_aliases( - call_type: CallTypesLiteral, - payload_key: Literal["messages", "input", "prompt"], - is_valid: bool, - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(litellm, "llm_guard_mode", "all") - llm_guard: Final = _ENTERPRISE_LLMGuard( - mock_testing=True, - mock_redacted_text={ - "sanitized_prompt": "email: [REDACTED]", - "is_valid": is_valid, - }, - ) - user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-12345")) - data: Final = { - payload_key: [{"role": "user", "content": "email: person@example.com"}] - if payload_key == "messages" - else "email: person@example.com" - } - - if not is_valid: - with pytest.raises(HTTPException) as exc_info: - await llm_guard.async_moderation_hook( - data=data, user_api_key_dict=user_api_key_dict, call_type=call_type - ) - assert exc_info.value.status_code == 400 - assert exc_info.value.detail == {"error": "Violated content safety policy"} - return - - result: Final = await llm_guard.async_moderation_hook( - data=data, user_api_key_dict=user_api_key_dict, call_type=call_type - ) - assert result is data - assert data[payload_key] == ( - [{"role": "user", "content": "email: [REDACTED]"}] - if payload_key == "messages" - else "email: [REDACTED]" - ) - - -@pytest.mark.parametrize( - "call_type", - ( - "responses", - "aresponses", - "anthropic_messages", - "aanthropic_messages", - "aspeech", - "aimage_edit", - "pass_through_endpoint", - ), -) -@pytest.mark.asyncio -async def test_llm_guard_skips_unsupported_call_types( - call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "llm_guard_mode", "all") - llm_guard: Final = _ENTERPRISE_LLMGuard( - mock_testing=True, - mock_redacted_text={"is_valid": False}, - ) - data: Final = {"messages": [{"role": "user", "content": "unchanged"}]} - result: Final = await llm_guard.async_moderation_hook( - data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type - ) - assert result is data - assert data == {"messages": [{"role": "user", "content": "unchanged"}]} - - @pytest.mark.asyncio async def test_llm_guard_error_raising(): """ diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py new file mode 100644 index 00000000000..dcb14e176e9 --- /dev/null +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py @@ -0,0 +1,97 @@ +from typing import Final, Literal + +import pytest +from fastapi import HTTPException +from litellm_enterprise.enterprise_callbacks.llm_guard import _ENTERPRISE_LLMGuard + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import hash_token +from litellm.types.utils import CallTypesLiteral + + +@pytest.mark.parametrize( + "call_type, payload_key", + ( + ("completion", "messages"), + ("acompletion", "messages"), + ("text_completion", "prompt"), + ("atext_completion", "prompt"), + ("embeddings", "input"), + ("embedding", "input"), + ("aembedding", "input"), + ("moderation", "input"), + ("amoderation", "input"), + ("image_generation", "prompt"), + ("aimage_generation", "prompt"), + ("audio_transcription", "prompt"), + ("transcription", "prompt"), + ("atranscription", "prompt"), + ), +) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_call_type_aliases( + call_type: CallTypesLiteral, + payload_key: Literal["messages", "input", "prompt"], + is_valid: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={ + "sanitized_prompt": "email: [REDACTED]", + "is_valid": is_valid, + }, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-12345")) + data: Final = { + payload_key: [{"role": "user", "content": "email: person@example.com"}] + if payload_key == "messages" + else "email: person@example.com" + } + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=user_api_key_dict, call_type=call_type) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated content safety policy"} + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=call_type + ) + assert result is data + assert data[payload_key] == ( + [{"role": "user", "content": "email: [REDACTED]"}] if payload_key == "messages" else "email: [REDACTED]" + ) + + +@pytest.mark.parametrize( + "call_type", + ( + "responses", + "aresponses", + "anthropic_messages", + "aanthropic_messages", + "aspeech", + "aimage_edit", + "pass_through_endpoint", + ), +) +@pytest.mark.asyncio +async def test_llm_guard_skips_unsupported_call_types( + call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"is_valid": False}, + ) + data: Final = {"messages": [{"role": "user", "content": "unchanged"}]} + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data == {"messages": [{"role": "user", "content": "unchanged"}]} From 78a29ae08f0c36b05163c5c475cff39f0eb33843 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 08:00:24 +0000 Subject: [PATCH 3/5] fix(llmguard): scan list valued completion prompts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../enterprise_callbacks/llm_guard.py | 18 ++++++------- .../enterprise_callbacks/test_llm_guard.py | 26 +++++++++++++++++++ 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index 9c8537e6820..7338352106a 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -177,12 +177,12 @@ class _ENTERPRISE_LLMGuard(CustomLogger): input_ = data.get("input") if input_ is not None: - data["input"] = await self._moderate_input(input_) + data["input"] = await self._moderate_text_or_list(input_) return data prompt = data.get("prompt") - if isinstance(prompt, str): - data["prompt"] = await self.moderation_check(text=prompt) + if prompt is not None: + data["prompt"] = await self._moderate_text_or_list(prompt) return data async def _moderate_message(self, message: dict) -> dict: @@ -205,17 +205,17 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return {**part, "text": await self.moderation_check(text=part["text"])} return part - async def _moderate_input(self, input_: object) -> object: - if isinstance(input_, str): - return await self.moderation_check(text=input_) - if isinstance(input_, list): + async def _moderate_text_or_list(self, value: object) -> object: + if isinstance(value, str): + return await self.moderation_check(text=value) + if isinstance(value, list): return [ await self.moderation_check(text=item) if isinstance(item, str) else item - for item in input_ + for item in value ] - return input_ + return value async def async_post_call_streaming_hook( self, user_api_key_dict: UserAPIKeyAuth, response: str diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py index dcb14e176e9..ef2aa96c36f 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py @@ -68,6 +68,32 @@ async def test_llm_guard_call_type_aliases( ) +@pytest.mark.parametrize("call_type", ("text_completion", "atext_completion")) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_scans_list_prompt( + call_type: CallTypesLiteral, is_valid: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": is_valid}, + ) + data: Final = {"prompt": ["email: person@example.com", "say ok", [1, 2, 3]]} + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type) + assert exc_info.value.status_code == 400 + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["prompt"] == ["[REDACTED]", "[REDACTED]", [1, 2, 3]] + + @pytest.mark.parametrize( "call_type", ( From 012d82d85ddac98cb81931e100159968f1eb8e3d Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 16:30:11 +0000 Subject: [PATCH 4/5] fix(llmguard): scan input and prompt even when messages is present Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../enterprise_callbacks/llm_guard.py | 2 -- .../enterprise_callbacks/test_llm_guard.py | 34 ++++++++++++++++++- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index 7338352106a..1559fff291c 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -173,12 +173,10 @@ class _ENTERPRISE_LLMGuard(CustomLogger): *(self._moderate_message(message) for message in messages) ) ) - return data input_ = data.get("input") if input_ is not None: data["input"] = await self._moderate_text_or_list(input_) - return data prompt = data.get("prompt") if prompt is not None: diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py index ef2aa96c36f..4bb663b3bf0 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py @@ -1,8 +1,8 @@ from typing import Final, Literal import pytest -from fastapi import HTTPException from litellm_enterprise.enterprise_callbacks.llm_guard import _ENTERPRISE_LLMGuard +from starlette.exceptions import HTTPException import litellm from litellm.proxy._types import UserAPIKeyAuth @@ -94,6 +94,38 @@ async def test_llm_guard_scans_list_prompt( assert data["prompt"] == ["[REDACTED]", "[REDACTED]", [1, 2, 3]] +@pytest.mark.parametrize("call_type", ("aembedding", "atext_completion")) +@pytest.mark.parametrize("is_valid", (True, False)) +@pytest.mark.asyncio +async def test_llm_guard_scans_input_and_prompt_alongside_messages( + call_type: CallTypesLiteral, is_valid: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": is_valid}, + ) + data: Final = { + "messages": [], + "input": "email: person@example.com", + "prompt": ["say ok"], + } + + if not is_valid: + with pytest.raises(HTTPException) as exc_info: + await llm_guard.async_moderation_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type) + assert exc_info.value.status_code == 400 + return + + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["messages"] == [] + assert data["input"] == "[REDACTED]" + assert data["prompt"] == ["[REDACTED]"] + + @pytest.mark.parametrize( "call_type", ( From c38dda2b2f47cd7e1056f077b8943f598c26cd21 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 17:57:47 +0000 Subject: [PATCH 5/5] fix(llmguard): drop call types the proxy never routes through moderation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../enterprise_callbacks/llm_guard.py | 5 ---- .../enterprise_callbacks/test_llm_guard.py | 23 +++++++++++++++---- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index 1559fff291c..3422e8969b0 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -147,11 +147,6 @@ class _ENTERPRISE_LLMGuard(CustomLogger): "aembedding", "image_generation", "aimage_generation", - "moderation", - "amoderation", - "audio_transcription", - "transcription", - "atranscription", ) if call_type not in accepted_call_types: self.print_verbose( diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py index 4bb663b3bf0..5695b184479 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py @@ -20,13 +20,8 @@ from litellm.types.utils import CallTypesLiteral ("embeddings", "input"), ("embedding", "input"), ("aembedding", "input"), - ("moderation", "input"), - ("amoderation", "input"), ("image_generation", "prompt"), ("aimage_generation", "prompt"), - ("audio_transcription", "prompt"), - ("transcription", "prompt"), - ("atranscription", "prompt"), ), ) @pytest.mark.parametrize("is_valid", (True, False)) @@ -68,6 +63,24 @@ async def test_llm_guard_call_type_aliases( ) +@pytest.mark.parametrize("call_type", ("amoderation", "atranscription", "aresponses", "aanthropic_messages")) +@pytest.mark.asyncio +async def test_llm_guard_ignores_call_types_the_proxy_never_moderates( + call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "llm_guard_mode", "all") + llm_guard: Final = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={"sanitized_prompt": "[REDACTED]", "is_valid": False}, + ) + data: Final = {"input": "email: person@example.com"} + result: Final = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ) + assert result is data + assert data["input"] == "email: person@example.com" + + @pytest.mark.parametrize("call_type", ("text_completion", "atext_completion")) @pytest.mark.parametrize("is_valid", (True, False)) @pytest.mark.asyncio