Add in-enclave moderation of image requests via OpenAI omni-moderation

CLAUDE.md

@@ -182,6 +182,37 @@ to the relay or the gateway operator. Points to keep in mind:
 - Exa's self-reported `costDollars` is logged for margin reconciliation only;
   settlement never depends on it.
 
+### Content Moderation
+
+Image requests on `/v1/chat/completions` — image generation, image editing,
+and inline-image chat models — are scored against OpenAI's free
+`omni-moderation-latest` endpoint before any provider is called
+(`moderation.py`). The check covers the newest user turn: its prompt text plus
+any attached images. **Plain text chat is not moderated**; widening scope
+there is a deliberate future change (the `should_moderate_model` predicate is
+the one gate to widen — everything downstream keys off the flag headers).
+Points to keep in mind:
+
+- **Fail-open**: no OpenAI key or a moderation outage means requests proceed
+  unscored (`checked: false`); a positive verdict always comes from a real
+  moderation response. `/health` reports `moderation_enabled`.
+- **Blocking**: a request flagged for a category in
+  `moderation.BLOCKED_CATEGORIES` (default: `sexual/minors` only) is refused
+  with HTTP 451 + `code: "moderation_blocked"` and never reaches a provider.
+  All other flagged categories are reported but still served.
+- **Response surface**: the full verdict (flagged/blocked/categories/scores)
+  rides inside the sealed response body under the `moderation` key —
+  non-streaming responses and the final SSE frame alike — outside the signed
+  output hash, exactly like `images` and `usage`.
+- **Relay signal**: flagged requests additionally carry content-free
+  `X-Moderation-Flagged` / `X-Moderation-Categories` / `X-Moderation-Blocked`
+  outer headers (forwarded through the OHTTP path) so the relay can run its
+  per-user strike/blacklist policy. Clean traffic carries none of these — it
+  is byte-identical to before.
+- **Billing**: the moderation call is free and adds no cost block changes;
+  blocked (451) requests produce no `opengradient` block and are never
+  settled.
+
 ## Verification Examples
 
 - `examples/verify_attestation.py` — Validates AWS Nitro attestation documents against the root CA

tee_gateway/__main__.py

@@ -22,6 +22,7 @@
     DEFAULT_HEARTBEAT_INTERVAL,
 )
 from tee_gateway.llm_backend import get_provider_config, set_provider_config
+from tee_gateway.moderation import moderation_available
 from tee_gateway.web_search import web_search_available
 from tee_gateway.heartbeat import create_heartbeat_service
 from tee_gateway.controllers.ohttp_controller import (
@@ -534,6 +535,7 @@ def _set(val: str | None) -> str:
             "providers_initialized": providers_set,
             "heartbeat_enabled": heartbeat_config is not None,
             "web_search_enabled": bool(provider_config.exa_api_key),
+            "moderation_enabled": bool(provider_config.openai_api_key),
         }
     ), 200
 
@@ -552,6 +554,9 @@ def health():
         # Not a provider capability — the search endpoint has no model — so it
         # is reported separately from `providers`.
         "web_search_enabled": web_search_available(),
+        # Whether chat prompts are being scored against OpenAI's moderation
+        # endpoint (requires the OpenAI key; fail-open when unavailable).
+        "moderation_enabled": moderation_available(),
         "facilitator_url": _active_facilitator_url,
         "price_feed": _price_feed.get_status(),
     }, 200

tee_gateway/controllers/chat_controller.py

@@ -39,6 +39,11 @@
     create_image_generation_streaming_response,
 )
 from tee_gateway.model_registry import get_model_config
+from tee_gateway.moderation import (
+    ModerationOutcome,
+    moderate_messages,
+    should_moderate_model,
+)
 from tee_gateway.pricing import compute_session_cost
 
 logger = logging.getLogger(__name__)
@@ -97,10 +102,37 @@ def create_chat_completion(body):
     except AttachmentValidationError as e:
         return {"error": "Invalid attachment", "message": str(e)}, 400
 
+    # Score the newest user turn of image requests before any provider work
+    # (text chat is not moderated — see moderation.should_moderate_model).
+    # Fail-open: an unavailable moderation endpoint yields an unchecked
+    # outcome, never a refusal. Only BLOCKED_CATEGORIES verdicts stop the
+    # request here; other flags ride along on the response for the
+    # relay/client to act on.
+    moderation = (
+        moderate_messages(chat_request.messages)
+        if should_moderate_model(chat_request.model)
+        else None
+    )
+    if moderation is not None and moderation.blocked:
+        return (
+            {
+                "error": "Content policy violation",
+                "message": (
+                    "This request was flagged by content moderation for a "
+                    "prohibited category and was not forwarded to the model "
+                    "provider."
+                ),
+                "code": "moderation_blocked",
+                "moderation": moderation.to_response_dict(),
+            },
+            451,
+            moderation.headers(),
+        )
+
     if chat_request.stream:
-        return _create_streaming_response(chat_request)
+        return _create_streaming_response(chat_request, moderation)
     else:
-        return _create_non_streaming_response(chat_request)
+        return _create_non_streaming_response(chat_request, moderation)
 
 
 def _build_tools_list(chat_request: CreateChatCompletionRequest) -> list:
@@ -222,7 +254,10 @@ def _messages_contain_json_word(messages: list) -> bool:
     return False
 
 
-def _create_non_streaming_response(chat_request: CreateChatCompletionRequest):
+def _create_non_streaming_response(
+    chat_request: CreateChatCompletionRequest,
+    moderation: ModerationOutcome | None = None,
+):
     """Handle non-streaming chat completion via direct LangChain call."""
     try:
         logger.info("=" * 80)
@@ -241,7 +276,9 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest):
         # surfaced through this same endpoint, returning images out-of-band just
         # like Gemini's inline-image models.
         if cfg.image_generation:
-            return create_image_generation_response(chat_request, request_bytes)
+            return create_image_generation_response(
+                chat_request, request_bytes, moderation
+            )
 
         # Build the tools list first: some OpenAI models (gpt-5.6 family) must be
         # constructed against the Responses API when function tools are bound.
@@ -376,8 +413,16 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest):
                 provider,
             )
 
+        # The moderation verdict rides inside the (sealed) response body like
+        # images/usage — not part of the output hash. Flagged requests also get
+        # content-free X-Moderation-* headers for the relay's strike policy.
+        if moderation is not None and moderation.checked:
+            openai_response["moderation"] = moderation.to_response_dict()
+
         # Validate schema (the extra tee_* fields are preserved by returning dict directly)
         CreateChatCompletionResponse.from_dict(openai_response)
+        if moderation is not None and moderation.flagged:
+            return openai_response, 200, moderation.headers()
         return openai_response
 
     except Exception as e:
@@ -388,7 +433,10 @@ def _create_non_streaming_response(chat_request: CreateChatCompletionRequest):
         }, 500
 
 
-def _create_streaming_response(chat_request: CreateChatCompletionRequest):
+def _create_streaming_response(
+    chat_request: CreateChatCompletionRequest,
+    moderation: ModerationOutcome | None = None,
+):
     """Handle streaming chat completion via direct LangChain call."""
     try:
         provider = get_provider_from_model(chat_request.model)
@@ -407,7 +455,7 @@ def _create_streaming_response(chat_request: CreateChatCompletionRequest):
         # images endpoint; handle them without building a chat model.
         if get_model_config(chat_request.model).image_generation:
             return create_image_generation_streaming_response(
-                chat_request, request_bytes
+                chat_request, request_bytes, moderation
             )
 
         # Build the tools list first: some OpenAI models (gpt-5.6 family) must be
@@ -765,6 +813,10 @@ def generate():
                 # are not part of the signed output hash.
                 if generated_images:
                     final_data["images"] = generated_images
+                # Moderation verdict rides the final frame the same way (the
+                # flag headers were already flushed with the response headers).
+                if moderation is not None and moderation.checked:
+                    final_data["moderation"] = moderation.to_response_dict()
 
                 logger.debug(
                     f"Response Final\n\tTEE Signature: {tee_signature}\n\tTEE request hash: {input_hash_hex}\n\tTEE output hash: {output_hash_hex}\n\tTEE timestamp: {timestamp}\n\tTEE ID: 0x{tee_keys.get_tee_id()}"
@@ -825,6 +877,7 @@ def generate():
             headers={
                 "Cache-Control": "no-cache",
                 "Connection": "keep-alive",
+                **(moderation.headers() if moderation is not None else {}),
             },
         )
 

tee_gateway/controllers/ohttp_controller.py

@@ -113,8 +113,17 @@
 }
 
 # Response headers we propagate from the inner /v1/chat/completions response
-# back through the relay to the client.
-_FORWARDED_HEADER_PREFIXES = ("x-payment", "x-upto", "x-settlement", "x-tee")
+# back through the relay to the client. x-moderation-* is a deliberate, narrow
+# exception to "the relay learns nothing": a content-free flagged/blocked bit
+# (plus category names) the relay needs to run its per-user strike/blacklist
+# policy. It is only emitted on flagged requests — clean traffic carries none.
+_FORWARDED_HEADER_PREFIXES = (
+    "x-payment",
+    "x-upto",
+    "x-settlement",
+    "x-tee",
+    "x-moderation",
+)
 _FORWARDED_HEADER_NAMES = ("www-authenticate",)
 
 

tee_gateway/image_generation.py

@@ -35,6 +35,7 @@
 from flask import Response
 
 from tee_gateway import llm_backend
+from tee_gateway.moderation import ModerationOutcome
 from tee_gateway.models.create_chat_completion_request import (
     CreateChatCompletionRequest,
 )
@@ -519,7 +520,9 @@ def _run_image_generation(
 
 
 def create_image_generation_response(
-    chat_request: CreateChatCompletionRequest, request_bytes: bytes
+    chat_request: CreateChatCompletionRequest,
+    request_bytes: bytes,
+    moderation: ModerationOutcome | None = None,
 ):
     """Non-streaming image generation via a provider's images endpoint.
 
@@ -547,13 +550,21 @@ def create_image_generation_response(
     }
     if result["opengradient"] is not None:
         openai_response["opengradient"] = result["opengradient"]
+    # Prompt moderation verdict (computed by the chat controller before the
+    # generation ran) rides inside the sealed body like the images themselves.
+    if moderation is not None and moderation.checked:
+        openai_response["moderation"] = moderation.to_response_dict()
 
     CreateChatCompletionResponse.from_dict(openai_response)
+    if moderation is not None and moderation.flagged:
+        return openai_response, 200, moderation.headers()
     return openai_response
 
 
 def create_image_generation_streaming_response(
-    chat_request: CreateChatCompletionRequest, request_bytes: bytes
+    chat_request: CreateChatCompletionRequest,
+    request_bytes: bytes,
+    moderation: ModerationOutcome | None = None,
 ):
     """Streaming image generation: image gen is not a token stream, so we invoke
     once and emit the result on the final SSE frame (mirrors the Gemini path)."""
@@ -577,6 +588,8 @@ def generate():
                 final_data["images"] = result["images"]
             if result["opengradient"] is not None:
                 final_data["opengradient"] = result["opengradient"]
+            if moderation is not None and moderation.checked:
+                final_data["moderation"] = moderation.to_response_dict()
 
             yield f"data: {json.dumps(final_data)}\n\n"
             yield "data: [DONE]\n\n"
@@ -594,5 +607,9 @@ def generate():
     return Response(
         generate(),
         mimetype="text/event-stream",
-        headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            **(moderation.headers() if moderation is not None else {}),
+        },
     )

tee_gateway/llm_backend.py

@@ -28,6 +28,7 @@
 
 from tee_gateway.config import ProviderConfig
 from tee_gateway.model_registry import get_model_config
+from tee_gateway.moderation import configure_moderation_client
 from tee_gateway.web_search import configure_exa_client
 
 logger = logging.getLogger(__name__)
@@ -136,6 +137,10 @@ def set_provider_config(config: ProviderConfig) -> None:
     # provider's native tool, so its client is built here alongside them.
     configure_exa_client(config.exa_api_key)
 
+    # Prompt moderation reuses the OpenAI key but gets its own client with a
+    # much tighter timeout — it runs synchronously in front of every chat call.
+    configure_moderation_client(config.openai_api_key)
+
     get_chat_model_cached.cache_clear()
     _provider_config = config
 

tee_gateway/models/create_chat_completion_response.py

@@ -11,6 +11,7 @@ def __init__(
         system_fingerprint=None,
         object=None,
         usage=None,
+        moderation=None,
     ):
         self.id = id
         self.choices = choices
@@ -20,6 +21,7 @@ def __init__(
         self.system_fingerprint = system_fingerprint
         self.object = object
         self.usage = usage
+        self.moderation = moderation
 
     @classmethod
     def from_dict(cls, dikt) -> "CreateChatCompletionResponse":
@@ -34,4 +36,5 @@ def from_dict(cls, dikt) -> "CreateChatCompletionResponse":
             system_fingerprint=dikt.get("system_fingerprint"),
             object=dikt.get("object"),
             usage=dikt.get("usage"),
+            moderation=dikt.get("moderation"),
         )