Underground-Digital/Workflow-Engine

1from abc import ABC2from enum import Enum3from typing import Optional4 5from pydantic import BaseModel, Field, field_validator6 7 8class PromptMessageRole(Enum):9    """10    Enum class for prompt message.11    """12 13    SYSTEM = "system"14    USER = "user"15    ASSISTANT = "assistant"16    TOOL = "tool"17 18    @classmethod19    def value_of(cls, value: str) -> "PromptMessageRole":20        """21        Get value of given mode.22 23        :param value: mode value24        :return: mode25        """26        for mode in cls:27            if mode.value == value:28                return mode29        raise ValueError(f"invalid prompt message type value {value}")30 31 32class PromptMessageTool(BaseModel):33    """34    Model class for prompt message tool.35    """36 37    name: str38    description: str39    parameters: dict40 41 42class PromptMessageFunction(BaseModel):43    """44    Model class for prompt message function.45    """46 47    type: str = "function"48    function: PromptMessageTool49 50 51class PromptMessageContentType(Enum):52    """53    Enum class for prompt message content type.54    """55 56    TEXT = "text"57    IMAGE = "image"58    AUDIO = "audio"59 60 61class PromptMessageContent(BaseModel):62    """63    Model class for prompt message content.64    """65 66    type: PromptMessageContentType67    data: str68 69 70class TextPromptMessageContent(PromptMessageContent):71    """72    Model class for text prompt message content.73    """74 75    type: PromptMessageContentType = PromptMessageContentType.TEXT76 77 78class AudioPromptMessageContent(PromptMessageContent):79    type: PromptMessageContentType = PromptMessageContentType.AUDIO80    data: str = Field(..., description="Base64 encoded audio data")81    format: str = Field(..., description="Audio format")82 83 84class ImagePromptMessageContent(PromptMessageContent):85    """86    Model class for image prompt message content.87    """88 89    class DETAIL(str, Enum):90        LOW = "low"91        HIGH = "high"92 93    type: PromptMessageContentType = PromptMessageContentType.IMAGE94    detail: DETAIL = DETAIL.LOW95 96 97class PromptMessage(ABC, BaseModel):98    """99    Model class for prompt message.100    """101 102    role: PromptMessageRole103    content: Optional[str | list[PromptMessageContent]] = None104    name: Optional[str] = None105 106    def is_empty(self) -> bool:107        """108        Check if prompt message is empty.109 110        :return: True if prompt message is empty, False otherwise111        """112        return not self.content113 114 115class UserPromptMessage(PromptMessage):116    """117    Model class for user prompt message.118    """119 120    role: PromptMessageRole = PromptMessageRole.USER121 122 123class AssistantPromptMessage(PromptMessage):124    """125    Model class for assistant prompt message.126    """127 128    class ToolCall(BaseModel):129        """130        Model class for assistant prompt message tool call.131        """132 133        class ToolCallFunction(BaseModel):134            """135            Model class for assistant prompt message tool call function.136            """137 138            name: str139            arguments: str140 141        id: str142        type: str143        function: ToolCallFunction144 145        @field_validator("id", mode="before")146        @classmethod147        def transform_id_to_str(cls, value) -> str:148            if not isinstance(value, str):149                return str(value)150            else:151                return value152 153    role: PromptMessageRole = PromptMessageRole.ASSISTANT154    tool_calls: list[ToolCall] = []155 156    def is_empty(self) -> bool:157        """158        Check if prompt message is empty.159 160        :return: True if prompt message is empty, False otherwise161        """162        if not super().is_empty() and not self.tool_calls:163            return False164 165        return True166 167 168class SystemPromptMessage(PromptMessage):169    """170    Model class for system prompt message.171    """172 173    role: PromptMessageRole = PromptMessageRole.SYSTEM174 175 176class ToolPromptMessage(PromptMessage):177    """178    Model class for tool prompt message.179    """180 181    role: PromptMessageRole = PromptMessageRole.TOOL182    tool_call_id: str183 184    def is_empty(self) -> bool:185        """186        Check if prompt message is empty.187 188        :return: True if prompt message is empty, False otherwise189        """190        if not super().is_empty() and not self.tool_call_id:191            return False192 193        return True194