| from typing import Optional, List |
| from backend.utils.sanitization import sanitize_user_input |
| from pydantic import BaseModel, Field, validator |
| from enum import Enum |
|
|
|
|
| class MessagePartType(str, Enum): |
| """Enumeration of supported message part types.""" |
| TEXT = "text" |
| STEP_START = "step-start" |
|
|
| class MessageRole(str, Enum): |
| """Enumeration of supported message roles.""" |
| USER = "user" |
| ASSISTANT = "assistant" |
| SYSTEM = "system" |
| |
| class MessagePart(BaseModel): |
| """Represents a single part of a message (text etc.).""" |
| type: MessagePartType |
| text: Optional[str] = Field(None, description="Text content for text-type parts") |
|
|
| @validator('text') |
| def validate_text_content(cls, v, values): |
| """Ensure text parts have non-empty text content.""" |
| if values.get('type') == MessagePartType.TEXT and (not v or not v.strip()): |
| raise ValueError("Text parts must have non-empty text content") |
| return v |
|
|
| class Config: |
| use_enum_values = True |
| |
|
|
| class Message(BaseModel): |
| """Represents a single message in a conversation.""" |
| id: str = Field(..., description="Unique identifier for the message part") |
| role: MessageRole |
| parts: List[MessagePart] = Field(..., min_items=1, description="Message content parts") |
| |
| @validator('parts') |
| def validate_parts_not_empty(cls, v): |
| """Ensure message has at least one valid part.""" |
| if not v: |
| raise ValueError("Message must contain at least one part") |
| return v |
|
|
| @validator('id') |
| def validate_id_format(cls, v): |
| """Ensure ID is non-empty string.""" |
| if not v or not v.strip(): |
| raise ValueError("Message ID cannot be empty") |
| return v.strip() |
|
|
| class Config: |
| use_enum_values = True |
|
|
| class ChatMessage(BaseModel): |
| """Represents a complete chat conversation.""" |
| id: str = Field(..., description="Unique identifier for the chat") |
| messages: List[Message] = Field(..., min_items=1, description="List of messages in the conversation") |
| trigger: Optional[str] = Field(None, description="Optional trigger that initiated the chat") |
|
|
| @validator('id') |
| def validate_chat_id(cls, v): |
| """Ensure chat ID is non-empty string.""" |
| if not v or not v.strip(): |
| raise ValueError("Chat ID cannot be empty") |
| return v.strip() |
|
|
| @validator('messages') |
| def validate_messages_structure(cls, v): |
| """Comprehensive validation of messages list.""" |
| if not v: |
| raise ValueError("Chat must contain at least one message") |
| |
| |
| message_ids = [msg.id for msg in v] |
| if len(message_ids) != len(set(message_ids)): |
| raise ValueError("All message IDs must be unique within a chat") |
| |
| |
| user_messages = [msg for msg in v if msg.role == MessageRole.USER] |
| if not user_messages: |
| raise ValueError("Chat must contain at least one user message") |
| return v |
|
|
| def get_latest_message(self) -> Optional[Message]: |
| """Get the most recent sanitized message in the chat.""" |
| latest = self.messages[-1] if self.messages else None |
| if latest and latest.parts: |
| for part in latest.parts: |
| if part.type == MessagePartType.TEXT and part.text: |
| part.text = sanitize_user_input(part.text) |
| return latest |
|
|
| class Config: |
| use_enum_values = True |
