mirror of https://github.com/kortix-ai/suna.git
89 lines
2.0 KiB
Python
89 lines
2.0 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Optional, Union, Dict, Any
|
|
from enum import Enum
|
|
|
|
|
|
class Role(str, Enum):
|
|
USER = "user"
|
|
ASSISTANT = "assistant"
|
|
SYSTEM = "system"
|
|
|
|
|
|
@dataclass
|
|
class ContentObject:
|
|
role: Role # "user" | "assistant" | "system"
|
|
content: Optional[str]
|
|
tool_calls: Optional[Any] = None # Can be made stricter if needed
|
|
|
|
|
|
class MessageType(str, Enum):
|
|
USER = "user"
|
|
ASSISTANT = "assistant"
|
|
TOOL = "tool"
|
|
STATUS = "status"
|
|
ASSISTANT_RESPONSE_END = "assistant_response_end"
|
|
|
|
|
|
@dataclass
|
|
class BaseMessage:
|
|
message_id: str
|
|
thread_id: str
|
|
type: MessageType # "user" | "assistant" | "tool" | "status" | "assistant_response_end"
|
|
is_llm_message: bool
|
|
metadata: Any
|
|
created_at: str
|
|
updated_at: str
|
|
|
|
|
|
@dataclass
|
|
class UserMessage(BaseMessage):
|
|
type: MessageType = field(init=False, default=MessageType.USER)
|
|
content: str # JSON string of ContentObject
|
|
|
|
|
|
@dataclass
|
|
class AssistantMessage(BaseMessage):
|
|
type: MessageType = field(init=False, default=MessageType.ASSISTANT)
|
|
content: ContentObject
|
|
|
|
|
|
@dataclass
|
|
class ToolResultMessage(BaseMessage):
|
|
type: MessageType = field(init=False, default=MessageType.TOOL)
|
|
content: Dict[
|
|
str, Union[str, Any]
|
|
] # role: "user", content: JSON string of ToolExecutionResult
|
|
|
|
|
|
@dataclass
|
|
class StatusMessage(BaseMessage):
|
|
type: MessageType = field(init=False, default=MessageType.STATUS)
|
|
content: Dict[str, Any] # status_type and other fields
|
|
|
|
|
|
@dataclass
|
|
class AssistantResponseEndMessage(BaseMessage):
|
|
type: MessageType = field(init=False, default=MessageType.ASSISTANT_RESPONSE_END)
|
|
content: Dict[str, Any] # model, usage, etc.
|
|
|
|
|
|
ChatMessage = Union[
|
|
UserMessage,
|
|
AssistantMessage,
|
|
ToolResultMessage,
|
|
StatusMessage,
|
|
AssistantResponseEndMessage,
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class AgentRun:
|
|
id: str
|
|
thread_id: str
|
|
status: str
|
|
started_at: Optional[str]
|
|
completed_at: Optional[str]
|
|
error: Any
|
|
created_at: str
|
|
updated_at: str
|