more xfail

tests/test_litellm_rust/integrations.py

@@ -15,8 +15,12 @@ import litellm
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig
 from litellm.integrations.prometheus import PrometheusLogger
+from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import (
+    AzureContentSafetyTextModerationGuardrail,
+)
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import GenericGuardrailAPIInputs
+from tests.test_litellm_rust.callback_recorder import drain_logging
 from tests.test_litellm_rust.contracts import (
     MESSAGES,
     MESSAGES_MODEL,
@@ -26,7 +30,6 @@ from tests.test_litellm_rust.contracts import (
     call_native_ocr,
 )
 from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.callback_recorder import drain_logging
 
 RouteName = Literal["ocr-sync", "ocr-async", "messages", "messages-stream"]
 GuardrailObservation = tuple[Literal["request", "response"], tuple[str, ...]]
@@ -118,12 +121,34 @@ MESSAGES_STREAM: Final = Route(
 ALL_ROUTES: Final = (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE, MESSAGES_STREAM)
 ASYNC_ROUTES: Final = tuple(route for route in ALL_ROUTES if route.fires_async_hooks)
 NON_STREAM_ASYNC_ROUTES: Final = (OCR_ASYNC, MESSAGES_ROUTE)
+AZURE_MODERATION_ALLOW_RESPONSE: Final = {
+    "blocklistsMatch": [],
+    "categoriesAnalysis": [
+        {"category": "Hate", "severity": 0},
+        {"category": "Sexual", "severity": 0},
+        {"category": "SelfHarm", "severity": 0},
+        {"category": "Violence", "severity": 0},
+    ],
+}
+AZURE_MODERATION_BLOCK_RESPONSE: Final = {
+    **AZURE_MODERATION_ALLOW_RESPONSE,
+    "categoriesAnalysis": [{"category": "Violence", "severity": 6}],
+}
 
 
 def route_id(route: Route) -> str:
     return route.name
 
 
+def azure_text_moderation(server: RecordingServer) -> AzureContentSafetyTextModerationGuardrail:
+    return AzureContentSafetyTextModerationGuardrail(
+        guardrail_name="azure-text-review",
+        api_key="test-azure-key",
+        api_base=server.base_url,
+        event_hook=GuardrailEventHooks.post_call,
+    )
+
+
 @pytest.fixture
 def provider(recording_server: RecordingServer, route: Route) -> RecordingServer:
     recording_server.default_response = ResponseSpec(body=route.provider_response)

tests/test_litellm_rust/test_integrations.py

@@ -19,6 +19,8 @@ from tests.test_litellm_rust.callback_recorder import RecordingLogger, drain_log
 from tests.test_litellm_rust.integrations import (
     ALL_ROUTES,
     ASYNC_ROUTES,
+    AZURE_MODERATION_ALLOW_RESPONSE,
+    AZURE_MODERATION_BLOCK_RESPONSE,
     MESSAGES_ROUTE,
     NON_STREAM_ASYNC_ROUTES,
     OCR_ASYNC,
@@ -27,6 +29,7 @@ from tests.test_litellm_rust.integrations import (
     RecordingGuardrail,
     ReviewGuardrail,
     Route,
+    azure_text_moderation,
     metric_value,
     route_id,
 )
@@ -65,7 +68,7 @@ async def test_generic_api_logger_exports_success_over_http(route: Route, provid
         pytest.param(
             MESSAGES_ROUTE,
             marks=pytest.mark.xfail(
-                reason="Rust Messages retries after a committed provider failure",
+                reason="Rust Messages retries HTTP 500 despite num_retries=0",
                 strict=True,
             ),
         ),
@@ -195,6 +198,64 @@ async def test_content_filter_post_call_blocks_provider_response(route: Route, p
     assert blocked.value.status_code == 400
 
 
+@pytest.mark.asyncio
+@pytest.mark.xfail(
+    strict=True,
+    raises=AssertionError,
+    reason="Azure text moderation only scans ModelResponse and skips native Messages responses",
+)
+async def test_azure_text_moderation_allows_messages_response_over_http(
+    recording_server: RecordingServer, otel: OtelHarness
+) -> None:
+    recording_server.expected_requests = None
+    recording_server.default_response = ResponseSpec(body=AZURE_MODERATION_ALLOW_RESPONSE)
+    recording_server.enqueue(ResponseSpec(body=MESSAGES_ROUTE.provider_response))
+    guardrail: Final = azure_text_moderation(recording_server)
+    recorder: Final = RecordingLogger()
+    litellm.callbacks.append(guardrail)
+
+    response: Final = await MESSAGES_ROUTE.invoke(
+        recording_server,
+        callbacks=[otel.logger, recorder],
+        guardrails=[guardrail.guardrail_name],
+    )
+    await recorder.wait_for_async("async_log_success_event")
+
+    assert len(recording_server.requests) == 2
+    moderation_request: Final = recording_server.requests[1]
+    assert moderation_request.path == "/contentsafety/text:analyze?api-version=2024-09-01"
+    assert moderation_request.headers["ocp-apim-subscription-key"] == "test-azure-key"
+    assert moderation_request.body == {
+        "text": MESSAGES_ROUTE.response_text,
+        "categories": ["Hate", "Sexual", "SelfHarm", "Violence"],
+        "blocklistNames": None,
+        "haltOnBlocklistHit": False,
+        "outputType": "FourSeverityLevels",
+    }
+    assert response["content"][0]["text"] == MESSAGES_ROUTE.response_text
+    assert len(await otel.wait_for_spans()) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.xfail(
+    strict=True,
+    raises=pytest.fail.Exception,
+    reason="Azure text moderation only scans ModelResponse and skips native Messages responses",
+)
+async def test_azure_text_moderation_blocks_messages_response_over_http(recording_server: RecordingServer) -> None:
+    recording_server.expected_requests = None
+    recording_server.default_response = ResponseSpec(body=AZURE_MODERATION_BLOCK_RESPONSE)
+    recording_server.enqueue(ResponseSpec(body=MESSAGES_ROUTE.provider_response))
+    guardrail: Final = azure_text_moderation(recording_server)
+    litellm.callbacks.append(guardrail)
+
+    with pytest.raises(HTTPException, match="Violence crossed severity 2") as blocked:
+        await MESSAGES_ROUTE.invoke(recording_server, guardrails=[guardrail.guardrail_name])
+
+    assert blocked.value.status_code == 400
+    assert recording_server.requests[1].body["text"] == MESSAGES_ROUTE.response_text
+
+
 def test_prometheus_registry_restores_collectors_after_failure() -> None:
     registry: Final = CollectorRegistry()
     original: Final = Counter("original", "Original collector", registry=registry)

tests/test_litellm_rust/test_messages_callbacks.py

@@ -171,7 +171,7 @@ async def test_messages_logging_drain_waits_for_suspended_callback(messages_serv
 
 @pytest.mark.asyncio
 @pytest.mark.xfail(
-    reason="Rust Messages retries after a committed provider failure",
+    reason="Rust Messages sends two provider requests before invoking failure callbacks",
     strict=True,
 )
 async def test_messages_failure_callbacks_receive_original_provider_error(messages_server: RecordingServer) -> None: