fix(server): validate chat completion messages

Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com>

Author: Pouyanpi

nemoguardrails/rails/llm/llmrails.py

@@ -807,7 +807,7 @@ def _get_events_for_messages(self, messages: List[dict], state: Any):
                     events.append({"type": "ContextUpdate", "data": msg["content"]})
                 elif msg["role"] == "event":
                     events.append(msg["event"])
-                elif msg["role"] == "system":
+                elif msg["role"] in ("developer", "system"):
                     # Handle system messages - convert them to SystemMessage events
                     events.append({"type": "SystemMessage", "content": msg["content"]})
                 elif msg["role"] == "tool":
@@ -864,7 +864,7 @@ def _get_events_for_messages(self, messages: List[dict], state: Any):
                     events.append({"type": "ContextUpdate", "data": msg["content"]})
                 elif msg["role"] == "event":
                     events.append(msg["event"])
-                elif msg["role"] == "system":
+                elif msg["role"] in ("developer", "system"):
                     # Handle system messages - convert them to SystemMessage events
                     events.append({"type": "SystemMessage", "content": msg["content"]})
                 elif msg["role"] == "tool":

nemoguardrails/server/schemas/openai.py

@@ -19,7 +19,16 @@
 from typing import Annotated, Any, List, Literal, Optional, Union
 
 from openai.types.chat.chat_completion import ChatCompletion
-from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, ValidationInfo, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    ValidationError,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
 
 from nemoguardrails.rails.llm.options import GenerationOptions, RailType
 
@@ -43,29 +52,193 @@ class GuardrailsChatCompletion(ChatCompletion):
     guardrails: Optional[GuardrailsDataOutput] = Field(default=None, description="Guardrails specific output data.")
 
 
-class _OpenAIChatMessageSchema(BaseModel):
-    model_config = ConfigDict(extra="allow")
+class _OpenAIChatMessageBase(BaseModel):
+    model_config = ConfigDict(extra="forbid")
 
-    role: str
+
+class _OpenAIChatMessageRoleSchema(BaseModel):
+    role: Literal["developer", "system", "user", "assistant", "tool", "function", "context"]
+
+
+class _OpenAIPromptCacheBreakpointSchema(_OpenAIChatMessageBase):
+    mode: Literal["explicit"]
+
+
+class _OpenAIImageURLSchema(_OpenAIChatMessageBase):
+    url: str
+    detail: Optional[Literal["auto", "low", "high"]] = None
+
+
+class _OpenAITextContentPartSchema(_OpenAIChatMessageBase):
+    text: str
+    type: Literal["text"]
+    prompt_cache_breakpoint: Optional[_OpenAIPromptCacheBreakpointSchema] = None
+
+
+class _OpenAIImageContentPartSchema(_OpenAIChatMessageBase):
+    image_url: _OpenAIImageURLSchema
+    type: Literal["image_url"]
+    prompt_cache_breakpoint: Optional[_OpenAIPromptCacheBreakpointSchema] = None
+
+
+class _OpenAIFileSchema(_OpenAIChatMessageBase):
+    file_data: Optional[str] = None
+    file_id: Optional[str] = None
+    filename: Optional[str] = None
+
+
+class _OpenAIFileContentPartSchema(_OpenAIChatMessageBase):
+    file: _OpenAIFileSchema
+    type: Literal["file"]
+    prompt_cache_breakpoint: Optional[_OpenAIPromptCacheBreakpointSchema] = None
+
+
+class _OpenAIRefusalContentPartSchema(_OpenAIChatMessageBase):
+    refusal: str
+    type: Literal["refusal"]
+
+
+_OpenAITextContentSchema = Union[str, List[_OpenAITextContentPartSchema]]
+_OpenAIUserContentPartSchema = Union[
+    _OpenAITextContentPartSchema,
+    _OpenAIImageContentPartSchema,
+    _OpenAIFileContentPartSchema,
+]
+
+
+class _OpenAIDeveloperMessageSchema(_OpenAIChatMessageBase):
+    content: _OpenAITextContentSchema
+    role: Literal["developer"]
+    name: Optional[str] = None
+
+
+class _OpenAISystemMessageSchema(_OpenAIChatMessageBase):
+    content: _OpenAITextContentSchema
+    role: Literal["system"]
+    name: Optional[str] = None
+
+
+class _OpenAIUserMessageSchema(_OpenAIChatMessageBase):
+    content: Union[str, List[_OpenAIUserContentPartSchema]]
+    role: Literal["user"]
+    name: Optional[str] = None
+
+
+class _OpenAIFunctionCallSchema(_OpenAIChatMessageBase):
+    arguments: str
+    name: str
+
+
+class _OpenAIFunctionToolCallSchema(_OpenAIChatMessageBase):
+    id: str
+    function: _OpenAIFunctionCallSchema
+    type: Literal["function"]
+
+
+class _OpenAICustomToolSchema(_OpenAIChatMessageBase):
+    input: str
+    name: str
+
+
+class _OpenAICustomToolCallSchema(_OpenAIChatMessageBase):
+    id: str
+    custom: _OpenAICustomToolSchema
+    type: Literal["custom"]
+
+
+_OpenAIToolCallSchema = Union[_OpenAIFunctionToolCallSchema, _OpenAICustomToolCallSchema]
+
+
+class _OpenAIAssistantMessageSchema(_OpenAIChatMessageBase):
+    role: Literal["assistant"]
+    content: Optional[Union[str, List[Union[_OpenAITextContentPartSchema, _OpenAIRefusalContentPartSchema]]]] = None
+    function_call: Optional[_OpenAIFunctionCallSchema] = None
+    name: Optional[str] = None
+    tool_calls: Optional[List[_OpenAIToolCallSchema]] = None
+    refusal: Optional[str] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_content_or_call(cls, value: Any) -> Any:
+        if (
+            isinstance(value, dict)
+            and value.get("content") is None
+            and value.get("function_call") is None
+            and not value.get("tool_calls")
+            and value.get("refusal") is None
+        ):
+            raise ValidationError.from_exception_data(
+                cls.__name__,
+                [{"type": "missing", "loc": ("content",), "input": value}],
+            )
+        return value
+
+
+class _OpenAIToolMessageSchema(_OpenAIChatMessageBase):
+    content: _OpenAITextContentSchema
+    role: Literal["tool"]
+    tool_call_id: str
+
+
+class _OpenAIFunctionMessageSchema(_OpenAIChatMessageBase):
+    content: Optional[str]
+    name: str
+    role: Literal["function"]
+
+
+class _GuardrailsContextMessageSchema(_OpenAIChatMessageBase):
+    content: dict[str, Any]
+    role: Literal["context"]
+
+
+_OpenAIChatMessageInput = Union[
+    _OpenAIDeveloperMessageSchema,
+    _OpenAISystemMessageSchema,
+    _OpenAIUserMessageSchema,
+    _OpenAIAssistantMessageSchema,
+    _OpenAIToolMessageSchema,
+    _OpenAIFunctionMessageSchema,
+    _GuardrailsContextMessageSchema,
+]
+
+_OPENAI_CHAT_MESSAGE_SCHEMAS: dict[str, type[BaseModel]] = {
+    "developer": _OpenAIDeveloperMessageSchema,
+    "system": _OpenAISystemMessageSchema,
+    "user": _OpenAIUserMessageSchema,
+    "assistant": _OpenAIAssistantMessageSchema,
+    "tool": _OpenAIToolMessageSchema,
+    "function": _OpenAIFunctionMessageSchema,
+    "context": _GuardrailsContextMessageSchema,
+}
 
 
 def _validate_openai_chat_message(message: Any) -> Any:
-    _OpenAIChatMessageSchema.model_validate(message)
+    role = _OpenAIChatMessageRoleSchema.model_validate(message).role
+    _OPENAI_CHAT_MESSAGE_SCHEMAS[role].model_validate(message)
     return message
 
 
 OpenAIChatMessage = Annotated[
     dict[str, Any],
     BeforeValidator(
         _validate_openai_chat_message,
-        json_schema_input_type=_OpenAIChatMessageSchema,
+        json_schema_input_type=_OpenAIChatMessageInput,
     ),
 ]
 
 
 class OpenAIChatCompletionRequest(BaseModel):
     """Standard OpenAI chat completion request parameters."""
 
+    @model_validator(mode="before")
+    @classmethod
+    def reject_unsupported_audio(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            modalities = data.get("modalities")
+            if "audio" in data or (isinstance(modalities, list) and "audio" in modalities):
+                raise ValueError("Audio input and output are not supported.")
+        return data
+
     messages: Optional[List[OpenAIChatMessage]] = Field(
         default=None,
         description="The list of messages in the current conversation.",