Handle sync vs async and oai vs azure clients

Author: steven10a

src/guardrails/checks/text/moderation.py

@@ -27,14 +27,21 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 from enum import Enum
 from functools import cache
 from typing import Any
 
-from openai import AsyncOpenAI, NotFoundError
+from openai import AsyncOpenAI, NotFoundError, OpenAI
 from pydantic import BaseModel, ConfigDict, Field
 
+try:
+    from openai import AsyncAzureOpenAI, AzureOpenAI  # type: ignore
+except Exception:  # pragma: no cover - optional dependency
+    AsyncAzureOpenAI = object  # type: ignore
+    AzureOpenAI = object  # type: ignore
+
 from guardrails.registry import default_spec_registry
 from guardrails.spec import GuardrailSpecMetadata
 from guardrails.types import GuardrailResult
@@ -130,11 +137,14 @@ def _get_moderation_client() -> AsyncOpenAI:
     return AsyncOpenAI()
 
 
-async def _call_moderation_api(client: AsyncOpenAI, data: str) -> Any:
-    """Call the OpenAI moderation API.
+async def _call_moderation_api_async(
+    client: AsyncOpenAI | AsyncAzureOpenAI,
+    data: str,  # type: ignore
+) -> Any:
+    """Call the OpenAI moderation API asynchronously.
 
     Args:
-        client: The OpenAI client to use.
+        client: The async OpenAI or Azure OpenAI client to use.
         data: The text to analyze.
 
     Returns:
@@ -146,6 +156,22 @@ async def _call_moderation_api(client: AsyncOpenAI, data: str) -> Any:
     )
 
 
+def _call_moderation_api_sync(client: OpenAI | AzureOpenAI, data: str) -> Any:  # type: ignore
+    """Call the OpenAI moderation API synchronously.
+
+    Args:
+        client: The sync OpenAI or Azure OpenAI client to use.
+        data: The text to analyze.
+
+    Returns:
+        The moderation API response.
+    """
+    return client.moderations.create(
+        model="omni-moderation-latest",
+        input=data,
+    )
+
+
 async def moderation(
     ctx: Any,
     data: str,
@@ -169,8 +195,15 @@ async def moderation(
     client = getattr(ctx, "guardrail_llm", None) if ctx is not None else None
 
     if client is not None:
+        # Determine if client is async or sync and call appropriately
+        is_async_client = isinstance(client, AsyncOpenAI | AsyncAzureOpenAI)
+
         try:
-            resp = await _call_moderation_api(client, data)
+            if is_async_client:
+                resp = await _call_moderation_api_async(client, data)
+            else:
+                # Sync client - run in thread pool to avoid blocking event loop
+                resp = await asyncio.to_thread(_call_moderation_api_sync, client, data)
         except NotFoundError as e:
             # Moderation endpoint doesn't exist on this provider (e.g., third-party)
             # Fall back to the OpenAI client
@@ -179,11 +212,11 @@ async def moderation(
                 e,
             )
             client = _get_moderation_client()
-            resp = await _call_moderation_api(client, data)
+            resp = await _call_moderation_api_async(client, data)
     else:
         # No context client, use fallback
         client = _get_moderation_client()
-        resp = await _call_moderation_api(client, data)
+        resp = await _call_moderation_api_async(client, data)
     results = resp.results or []
     if not results:
         return GuardrailResult(

tests/unit/checks/test_moderation.py

@@ -134,3 +134,35 @@ async def raise_not_found(**_: Any) -> Any:
     # Verify the fallback client was used (not the third-party one)
     assert fallback_used is True  # noqa: S101
     assert result.tripwire_triggered is False  # noqa: S101
+
+
+@pytest.mark.asyncio
+async def test_moderation_uses_sync_context_client() -> None:
+    """Moderation should support synchronous OpenAI clients from context."""
+    from openai import OpenAI
+
+    # Track whether sync context client was used
+    sync_client_used = False
+
+    def track_sync_create(**_: Any) -> Any:
+        nonlocal sync_client_used
+        sync_client_used = True
+
+        class _Result:
+            def model_dump(self) -> dict[str, Any]:
+                return {"categories": {"hate": False, "violence": False}}
+
+        return SimpleNamespace(results=[_Result()])
+
+    # Create a sync context client
+    sync_client = OpenAI(api_key="test-sync-key", base_url="https://api.openai.com/v1")
+    sync_client.moderations = SimpleNamespace(create=track_sync_create)  # type: ignore[assignment]
+
+    ctx = SimpleNamespace(guardrail_llm=sync_client)
+
+    cfg = ModerationCfg(categories=[Category.HATE, Category.VIOLENCE])
+    result = await moderation(ctx, "test text", cfg)
+
+    # Verify the sync context client was used (via asyncio.to_thread)
+    assert sync_client_used is True  # noqa: S101
+    assert result.tripwire_triggered is False  # noqa: S101