jl-ecms-client 0.2.8__py3-none-any.whl

mirix/schemas/message.py

@@ -0,0 +1,1315 @@
+from __future__ import annotations
+
+import copy
+import json
+import uuid
+import warnings
+from collections import OrderedDict
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator
+
+from mirix.client.constants import (
+    DEFAULT_MESSAGE_TOOL,
+    DEFAULT_MESSAGE_TOOL_KWARG,
+    TOOL_CALL_ID_MAX_LEN,
+)
+from mirix.helpers.datetime_helpers import get_utc_time, is_utc_datetime
+from mirix.helpers.json_helpers import json_dumps
+from mirix.schemas.enums import MessageRole
+from mirix.schemas.mirix_base import OrmMetadataBase
+from mirix.schemas.mirix_message import (
+    AssistantMessage,
+    HiddenReasoningMessage,
+    MirixMessage,
+    ReasoningMessage,
+    SystemMessage,
+    ToolCall,
+    ToolCallMessage,
+    ToolReturnMessage,
+    UserMessage,
+)
+from mirix.schemas.mirix_message_content import (
+    CloudFileContent,
+    FileContent,
+    ImageContent,
+    MirixMessageContentUnion,
+    ReasoningContent,
+    RedactedReasoningContent,
+    TextContent,
+    get_mirix_message_content_union_str_json_schema,
+)
+from mirix.schemas.openai.openai import Function as OpenAIFunction
+from mirix.schemas.openai.openai import ToolCall as OpenAIToolCall
+from mirix.system import unpack_message
+from mirix.helpers.json_helpers import parse_json
+
+class BaseMessage(OrmMetadataBase):
+    __id_prefix__ = "message"
+
+
+class MessageCreate(BaseModel):
+    """Request to create a message"""
+
+    # In the simplified format, only allow simple roles
+    role: Literal[
+        MessageRole.user,
+        MessageRole.system,
+    ] = Field(..., description="The role of the participant.")
+    content: Union[str, List[MirixMessageContentUnion]] = Field(
+        ...,
+        description="The content of the message.",
+        json_schema_extra=get_mirix_message_content_union_str_json_schema(),
+    )
+    name: Optional[str] = Field(None, description="The name of the participant.")
+    otid: Optional[str] = Field(
+        None, description="The offline threading id associated with this message"
+    )
+    sender_id: Optional[str] = Field(
+        None,
+        description="The id of the sender of the message, can be an identity id or agent id",
+    )
+    group_id: Optional[str] = Field(
+        None, description="The multi-agent group that the message was sent in"
+    )
+    filter_tags: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Optional tags for filtering and categorizing this message and related memories"
+    )
+
+    def model_dump(self, to_orm: bool = False, **kwargs) -> Dict[str, Any]:
+        data = super().model_dump(**kwargs)
+        if to_orm and "content" in data:
+            if isinstance(data["content"], str):
+                data["content"] = [TextContent(text=data["content"])]
+        return data
+
+
+class MessageUpdate(BaseModel):
+    """Request to update a message"""
+
+    role: Optional[MessageRole] = Field(
+        None, description="The role of the participant."
+    )
+    content: Optional[Union[str, List[MirixMessageContentUnion]]] = Field(
+        None,
+        description="The content of the message.",
+        json_schema_extra=get_mirix_message_content_union_str_json_schema(),
+    )
+    # NOTE: probably doesn't make sense to allow remapping user_id or agent_id (vs creating a new message)
+    # user_id: Optional[str] = Field(None, description="The unique identifier of the user.")
+    # agent_id: Optional[str] = Field(None, description="The unique identifier of the agent.")
+    # NOTE: we probably shouldn't allow updating the model field, otherwise this loses meaning
+    # model: Optional[str] = Field(None, description="The model used to make the function call.")
+    name: Optional[str] = Field(None, description="The name of the participant.")
+    # NOTE: we probably shouldn't allow updating the created_at field, right?
+    # created_at: Optional[datetime] = Field(None, description="The time the message was created.")
+    tool_calls: Optional[List[OpenAIToolCall,]] = Field(
+        None, description="The list of tool calls requested."
+    )
+    tool_call_id: Optional[str] = Field(None, description="The id of the tool call.")
+
+    def model_dump(self, to_orm: bool = False, **kwargs) -> Dict[str, Any]:
+        data = super().model_dump(**kwargs)
+        if to_orm and "content" in data:
+            if isinstance(data["content"], str):
+                data["content"] = [TextContent(text=data["content"])]
+        return data
+
+
+class Message(BaseMessage):
+    """
+    Mirix's internal representation of a message. Includes methods to convert to/from LLM provider formats.
+
+    Attributes:
+        id (str): The unique identifier of the message.
+        role (MessageRole): The role of the participant.
+        text (str): The text of the message.
+        user_id (str): The unique identifier of the user.
+        agent_id (str): The unique identifier of the agent.
+        model (str): The model used to make the function call.
+        name (str): The name of the participant.
+        created_at (datetime): The time the message was created.
+        tool_calls (List[OpenAIToolCall,]): The list of tool calls requested.
+        tool_call_id (str): The id of the tool call.
+
+    """
+
+    id: str = BaseMessage.generate_id_field()
+    organization_id: Optional[str] = Field(
+        None, description="The unique identifier of the organization."
+    )
+    user_id: Optional[str] = Field(
+        None, description="The unique identifier of the user."
+    )
+    agent_id: Optional[str] = Field(
+        None, description="The unique identifier of the agent."
+    )
+    model: Optional[str] = Field(
+        None, description="The model used to make the function call."
+    )
+    # Basic OpenAI-style fields
+    role: MessageRole = Field(..., description="The role of the participant.")
+    content: Optional[List[MirixMessageContentUnion]] = Field(
+        None, description="The content of the message."
+    )
+    # NOTE: in OpenAI, this field is only used for roles 'user', 'assistant', and 'function' (now deprecated). 'tool' does not use it.
+    name: Optional[str] = Field(
+        None,
+        description="For role user/assistant: the (optional) name of the participant. For role tool/function: the name of the function called.",
+    )
+    tool_calls: Optional[List[OpenAIToolCall]] = Field(
+        None,
+        description="The list of tool calls requested. Only applicable for role assistant.",
+    )
+    tool_call_id: Optional[str] = Field(
+        None, description="The ID of the tool call. Only applicable for role tool."
+    )
+    # Extras
+    step_id: Optional[str] = Field(
+        None, description="The id of the step that this message was created in."
+    )
+    otid: Optional[str] = Field(
+        None, description="The offline threading id associated with this message"
+    )
+    tool_returns: Optional[List[ToolReturn]] = Field(
+        None, description="Tool execution return information for prior tool calls"
+    )
+    group_id: Optional[str] = Field(
+        None, description="The multi-agent group that the message was sent in"
+    )
+    sender_id: Optional[str] = Field(
+        None,
+        description="The id of the sender of the message, can be an identity id or agent id",
+    )
+    # This overrides the optional base orm schema, created_at MUST exist on all messages objects
+    created_at: datetime = Field(
+        default_factory=get_utc_time,
+        description="The timestamp when the object was created.",
+    )
+    
+    # NEW: Filter tags for flexible filtering and categorization
+    filter_tags: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Custom filter tags for filtering and categorization",
+        examples=[
+            {
+                "project_id": "proj-abc",
+                "session_id": "sess-xyz",
+                "tags": ["important", "work"],
+                "priority": "high"
+            }
+        ]
+    )
+
+    @field_validator("role")
+    @classmethod
+    def validate_role(cls, v: str) -> str:
+        roles = ["system", "assistant", "user", "tool"]
+        assert v in roles, f"Role must be one of {roles}"
+        return v
+
+    def to_json(self):
+        json_message = vars(self)
+        if json_message["tool_calls"] is not None:
+            json_message["tool_calls"] = [vars(tc) for tc in json_message["tool_calls"]]
+        # turn datetime to ISO format
+        # also if the created_at is missing a timezone, add UTC
+        if not is_utc_datetime(self.created_at):
+            self.created_at = self.created_at.replace(tzinfo=timezone.utc)
+        json_message["created_at"] = self.created_at.isoformat()
+        return json_message
+
+    @staticmethod
+    def generate_otid():
+        return str(uuid.uuid4())
+
+    @staticmethod
+    def to_mirix_messages_from_list(
+        messages: List[Message],
+        use_assistant_message: bool = True,
+        assistant_message_tool_name: str = DEFAULT_MESSAGE_TOOL,
+        assistant_message_tool_kwarg: str = DEFAULT_MESSAGE_TOOL_KWARG,
+        reverse: bool = True,
+    ) -> List[MirixMessage]:
+        if use_assistant_message:
+            message_ids_to_remove = []
+            assistant_messages_by_tool_call = {
+                tool_call.id: msg
+                for msg in messages
+                if msg.role == MessageRole.assistant and msg.tool_calls
+                for tool_call in msg.tool_calls
+            }
+            for message in messages:
+                if (
+                    message.role == MessageRole.tool
+                    and message.tool_call_id in assistant_messages_by_tool_call
+                    and assistant_messages_by_tool_call[message.tool_call_id].tool_calls
+                    and assistant_message_tool_name
+                    in [
+                        tool_call.function.name
+                        for tool_call in assistant_messages_by_tool_call[
+                            message.tool_call_id
+                        ].tool_calls
+                    ]
+                ):
+                    message_ids_to_remove.append(message.id)
+
+            messages = [msg for msg in messages if msg.id not in message_ids_to_remove]
+
+        # Convert messages to MirixMessages
+        return [
+            msg
+            for m in messages
+            for msg in m.to_mirix_message(
+                use_assistant_message=use_assistant_message,
+                assistant_message_tool_name=assistant_message_tool_name,
+                assistant_message_tool_kwarg=assistant_message_tool_kwarg,
+                reverse=reverse,
+            )
+        ]
+
+    def to_mirix_message(
+        self,
+        use_assistant_message: bool = False,
+        assistant_message_tool_name: str = DEFAULT_MESSAGE_TOOL,
+        assistant_message_tool_kwarg: str = DEFAULT_MESSAGE_TOOL_KWARG,
+        reverse: bool = True,
+    ) -> List[MirixMessage]:
+        """Convert message object (in DB format) to the style used by the original Mirix API"""
+        messages = []
+
+        if self.role == MessageRole.assistant:
+            # Handle reasoning
+            if self.content:
+                # Check for ReACT-style COT inside of TextContent
+                if len(self.content) == 1 and isinstance(self.content[0], TextContent):
+                    otid = Message.generate_otid_from_id(self.id, len(messages))
+                    messages.append(
+                        ReasoningMessage(
+                            id=self.id,
+                            date=self.created_at,
+                            reasoning=self.content[0].text,
+                            name=self.name,
+                            otid=otid,
+                            sender_id=self.sender_id,
+                        )
+                    )
+                # Otherwise, we may have a list of multiple types
+                else:
+                    # TODO we can probably collapse these two cases into a single loop
+                    for content_part in self.content:
+                        otid = Message.generate_otid_from_id(self.id, len(messages))
+                        if isinstance(content_part, TextContent):
+                            # COT
+                            messages.append(
+                                ReasoningMessage(
+                                    id=self.id,
+                                    date=self.created_at,
+                                    reasoning=content_part.text,
+                                    name=self.name,
+                                    otid=otid,
+                                    sender_id=self.sender_id,
+                                )
+                            )
+                        elif isinstance(content_part, ReasoningContent):
+                            # "native" COT
+                            messages.append(
+                                ReasoningMessage(
+                                    id=self.id,
+                                    date=self.created_at,
+                                    reasoning=content_part.reasoning,
+                                    source="reasoner_model",  # TODO do we want to tag like this?
+                                    signature=content_part.signature,
+                                    name=self.name,
+                                    otid=otid,
+                                )
+                            )
+                        elif isinstance(content_part, RedactedReasoningContent):
+                            # "native" redacted/hidden COT
+                            messages.append(
+                                HiddenReasoningMessage(
+                                    id=self.id,
+                                    date=self.created_at,
+                                    state="redacted",
+                                    hidden_reasoning=content_part.data,
+                                    name=self.name,
+                                    otid=otid,
+                                    sender_id=self.sender_id,
+                                )
+                            )
+                        else:
+                            warnings.warn(
+                                f"Unrecognized content part in assistant message: {content_part}"
+                            )
+
+            if self.tool_calls is not None:
+                # This is type FunctionCall
+                for tool_call in self.tool_calls:
+                    otid = Message.generate_otid_from_id(self.id, len(messages))
+                    # If we're supporting using assistant message,
+                    # then we want to treat certain function calls as a special case
+                    if (
+                        use_assistant_message
+                        and tool_call.function.name == assistant_message_tool_name
+                    ):
+                        # We need to unpack the actual message contents from the function call
+                        try:
+                            func_args = parse_json(tool_call.function.arguments)
+                            message_string = func_args[assistant_message_tool_kwarg]
+                        except KeyError:
+                            raise ValueError(
+                                f"Function call {tool_call.function.name} missing {assistant_message_tool_kwarg} argument"
+                            )
+                        messages.append(
+                            AssistantMessage(
+                                id=self.id,
+                                date=self.created_at,
+                                content=message_string,
+                                name=self.name,
+                                otid=otid,
+                                sender_id=self.sender_id,
+                            )
+                        )
+                    else:
+                        messages.append(
+                            ToolCallMessage(
+                                id=self.id,
+                                date=self.created_at,
+                                tool_call=ToolCall(
+                                    name=tool_call.function.name,
+                                    arguments=tool_call.function.arguments,
+                                    tool_call_id=tool_call.id,
+                                ),
+                                name=self.name,
+                                otid=otid,
+                                sender_id=self.sender_id,
+                            )
+                        )
+        elif self.role == MessageRole.tool:
+            # This is type ToolReturnMessage
+            # Try to interpret the function return, recall that this is how we packaged:
+            # def package_function_response(was_success, response_string, timestamp=None):
+            #     formatted_time = get_local_time() if timestamp is None else timestamp
+            #     packaged_message = {
+            #         "status": "OK" if was_success else "Failed",
+            #         "message": response_string,
+            #         "time": formatted_time,
+            #     }
+            if (
+                self.content
+                and len(self.content) == 1
+                and isinstance(self.content[0], TextContent)
+            ):
+                text_content = self.content[0].text
+            else:
+                raise ValueError(
+                    f"Invalid tool return (no text object on message): {self.content}"
+                )
+
+            try:
+                function_return = parse_json(text_content)
+                status = function_return["status"]
+                if status == "OK":
+                    status_enum = "success"
+                elif status == "Failed":
+                    status_enum = "error"
+                else:
+                    raise ValueError(f"Invalid status: {status}")
+            except json.JSONDecodeError:
+                raise ValueError(f"Failed to decode function return: {text_content}")
+            assert self.tool_call_id is not None
+            messages.append(
+                # TODO make sure this is what the API returns
+                # function_return may not match exactly...
+                ToolReturnMessage(
+                    id=self.id,
+                    date=self.created_at,
+                    tool_return=text_content,
+                    status=self.tool_returns[0].status
+                    if self.tool_returns
+                    else status_enum,
+                    tool_call_id=self.tool_call_id,
+                    stdout=self.tool_returns[0].stdout if self.tool_returns else None,
+                    stderr=self.tool_returns[0].stderr if self.tool_returns else None,
+                    name=self.name,
+                    otid=self.id.replace("message-", ""),
+                    sender_id=self.sender_id,
+                )
+            )
+        elif self.role == MessageRole.user:
+            # This is type UserMessage
+            if (
+                self.content
+                and len(self.content) == 1
+                and isinstance(self.content[0], TextContent)
+            ):
+                text_content = self.content[0].text
+            elif self.content and len(self.content) > 1:
+                text_content = ""
+                for content in self.content:
+                    if isinstance(content, TextContent):
+                        text_content += content.text
+                    elif isinstance(content, ImageContent):
+                        text_content += "<image>" + content.image_id + "</image>"
+                    elif isinstance(content, FileContent):
+                        text_content += "<file>" + content.file_id + "</file>"
+                    elif isinstance(content, CloudFileContent):
+                        text_content += (
+                            "<cloud_file>" + content.cloud_file_uri + "</cloud_file>"
+                        )
+            else:
+                raise ValueError(
+                    f"Invalid user message (no text object on message): {self.content}"
+                )
+
+            message_str = unpack_message(text_content)
+            messages.append(
+                UserMessage(
+                    id=self.id,
+                    date=self.created_at,
+                    content=message_str or text_content,
+                    name=self.name,
+                    otid=self.otid,
+                    sender_id=self.sender_id,
+                )
+            )
+        elif self.role == MessageRole.system:
+            # This is type SystemMessage
+            if (
+                self.content
+                and len(self.content) == 1
+                and isinstance(self.content[0], TextContent)
+            ):
+                text_content = self.content[0].text
+            else:
+                raise ValueError(
+                    f"Invalid system message (no text object on system): {self.content}"
+                )
+
+            messages.append(
+                SystemMessage(
+                    id=self.id,
+                    date=self.created_at,
+                    content=text_content,
+                    name=self.name,
+                    otid=self.otid,
+                    sender_id=self.sender_id,
+                )
+            )
+        else:
+            raise ValueError(self.role)
+
+        if reverse:
+            messages.reverse()
+
+        return messages
+
+    @staticmethod
+    def dict_to_message(
+        agent_id: str,
+        openai_message_dict: dict,
+        model: Optional[str] = None,  # model used to make function call
+        allow_functions_style: bool = False,  # allow deprecated functions style?
+        created_at: Optional[datetime] = None,
+        id: Optional[str] = None,
+        name: Optional[str] = None,
+        group_id: Optional[str] = None,
+        tool_returns: Optional[List[ToolReturn]] = None,
+    ):
+        """Convert a ChatCompletion message object into a Message object (synced to DB)"""
+        if not created_at:
+            # timestamp for creation
+            created_at = get_utc_time()
+
+        assert "role" in openai_message_dict, openai_message_dict
+        assert "content" in openai_message_dict, openai_message_dict
+
+        # TODO(caren) implicit support for only non-parts/list content types
+        if (
+            openai_message_dict["content"] is not None
+            and type(openai_message_dict["content"]) is not str
+        ):
+            raise ValueError(
+                f"Invalid content type: {type(openai_message_dict['content'])}"
+            )
+        content = (
+            [TextContent(text=openai_message_dict["content"])]
+            if openai_message_dict["content"]
+            else []
+        )
+
+        # TODO(caren) bad assumption here that "reasoning_content" always comes before "redacted_reasoning_content"
+        if (
+            "reasoning_content" in openai_message_dict
+            and openai_message_dict["reasoning_content"]
+        ):
+            content.append(
+                ReasoningContent(
+                    reasoning=openai_message_dict["reasoning_content"],
+                    is_native=True,
+                    signature=(
+                        openai_message_dict["reasoning_content_signature"]
+                        if openai_message_dict["reasoning_content_signature"]
+                        else None
+                    ),
+                ),
+            )
+        if (
+            "redacted_reasoning_content" in openai_message_dict
+            and openai_message_dict["redacted_reasoning_content"]
+        ):
+            content.append(
+                RedactedReasoningContent(
+                    data=openai_message_dict["redacted_reasoning_content"]
+                    if "redacted_reasoning_content" in openai_message_dict
+                    else None,
+                ),
+            )
+
+        # If we're going from deprecated function form
+        if openai_message_dict["role"] == "function":
+            if not allow_functions_style:
+                raise DeprecationWarning(openai_message_dict)
+            assert "tool_call_id" in openai_message_dict, openai_message_dict
+
+            # Convert from 'function' response to a 'tool' response
+            if id is not None:
+                return Message(
+                    agent_id=agent_id,
+                    model=model,
+                    # standard fields expected in an OpenAI ChatCompletion message object
+                    role=MessageRole.tool,  # NOTE
+                    content=content,
+                    name=name,
+                    tool_calls=openai_message_dict["tool_calls"]
+                    if "tool_calls" in openai_message_dict
+                    else None,
+                    tool_call_id=openai_message_dict["tool_call_id"]
+                    if "tool_call_id" in openai_message_dict
+                    else None,
+                    created_at=created_at,
+                    id=str(id),
+                    tool_returns=tool_returns,
+                    group_id=group_id,
+                )
+            else:
+                return Message(
+                    agent_id=agent_id,
+                    model=model,
+                    # standard fields expected in an OpenAI ChatCompletion message object
+                    role=MessageRole.tool,  # NOTE
+                    content=content,
+                    name=name,
+                    tool_calls=openai_message_dict["tool_calls"]
+                    if "tool_calls" in openai_message_dict
+                    else None,
+                    tool_call_id=openai_message_dict["tool_call_id"]
+                    if "tool_call_id" in openai_message_dict
+                    else None,
+                    created_at=created_at,
+                    tool_returns=tool_returns,
+                    group_id=group_id,
+                )
+
+        elif (
+            "function_call" in openai_message_dict
+            and openai_message_dict["function_call"] is not None
+        ):
+            if not allow_functions_style:
+                raise DeprecationWarning(openai_message_dict)
+            assert openai_message_dict["role"] == "assistant", openai_message_dict
+            assert "tool_call_id" in openai_message_dict, openai_message_dict
+
+            # Convert a function_call (from an assistant message) into a tool_call
+            # NOTE: this does not conventionally include a tool_call_id (ToolCall.id), it's on the caster to provide it
+            tool_calls = [
+                OpenAIToolCall(
+                    id=openai_message_dict[
+                        "tool_call_id"
+                    ],  # NOTE: unconventional source, not to spec
+                    type="function",
+                    function=OpenAIFunction(
+                        name=openai_message_dict["function_call"]["name"],
+                        arguments=openai_message_dict["function_call"]["arguments"],
+                    ),
+                )
+            ]
+
+            if id is not None:
+                return Message(
+                    agent_id=agent_id,
+                    model=model,
+                    # standard fields expected in an OpenAI ChatCompletion message object
+                    role=MessageRole(openai_message_dict["role"]),
+                    content=content,
+                    name=name,
+                    tool_calls=tool_calls,
+                    tool_call_id=None,  # NOTE: None, since this field is only non-null for role=='tool'
+                    created_at=created_at,
+                    id=str(id),
+                    tool_returns=tool_returns,
+                    group_id=group_id,
+                )
+            else:
+                return Message(
+                    agent_id=agent_id,
+                    model=model,
+                    # standard fields expected in an OpenAI ChatCompletion message object
+                    role=MessageRole(openai_message_dict["role"]),
+                    content=content,
+                    name=openai_message_dict["name"]
+                    if "name" in openai_message_dict
+                    else None,
+                    tool_calls=tool_calls,
+                    tool_call_id=None,  # NOTE: None, since this field is only non-null for role=='tool'
+                    created_at=created_at,
+                    tool_returns=tool_returns,
+                    group_id=group_id,
+                )
+
+        else:
+            # Basic sanity check
+            if openai_message_dict["role"] == "tool":
+                assert (
+                    "tool_call_id" in openai_message_dict
+                    and openai_message_dict["tool_call_id"] is not None
+                ), openai_message_dict
+            else:
+                if "tool_call_id" in openai_message_dict:
+                    assert openai_message_dict["tool_call_id"] is None, (
+                        openai_message_dict
+                    )
+
+            if (
+                "tool_calls" in openai_message_dict
+                and openai_message_dict["tool_calls"] is not None
+            ):
+                assert openai_message_dict["role"] == "assistant", openai_message_dict
+
+                tool_calls = [
+                    OpenAIToolCall(
+                        id=tool_call["id"],
+                        type=tool_call["type"],
+                        function=tool_call["function"],
+                    )
+                    for tool_call in openai_message_dict["tool_calls"]
+                ]
+            else:
+                tool_calls = None
+
+            # If we're going from tool-call style
+            if id is not None:
+                return Message(
+                    agent_id=agent_id,
+                    model=model,
+                    # standard fields expected in an OpenAI ChatCompletion message object
+                    role=MessageRole(openai_message_dict["role"]),
+                    content=content,
+                    name=openai_message_dict["name"]
+                    if "name" in openai_message_dict
+                    else name,
+                    tool_calls=tool_calls,
+                    tool_call_id=openai_message_dict["tool_call_id"]
+                    if "tool_call_id" in openai_message_dict
+                    else None,
+                    created_at=created_at,
+                    id=str(id),
+                    tool_returns=tool_returns,
+                    group_id=group_id,
+                )
+            else:
+                return Message(
+                    agent_id=agent_id,
+                    model=model,
+                    # standard fields expected in an OpenAI ChatCompletion message object
+                    role=MessageRole(openai_message_dict["role"]),
+                    content=content,
+                    name=openai_message_dict["name"]
+                    if "name" in openai_message_dict
+                    else name,
+                    tool_calls=tool_calls,
+                    tool_call_id=openai_message_dict["tool_call_id"]
+                    if "tool_call_id" in openai_message_dict
+                    else None,
+                    created_at=created_at,
+                    tool_returns=tool_returns,
+                    group_id=group_id,
+                )
+
+    def to_openai_dict_search_results(
+        self, max_tool_id_length: int = TOOL_CALL_ID_MAX_LEN
+    ) -> dict:
+        result_json = self.to_openai_dict()
+        search_result_json = {
+            "timestamp": self.created_at,
+            "message": {"content": result_json["content"], "role": result_json["role"]},
+        }
+        return search_result_json
+
+    def to_openai_dict(
+        self,
+        max_tool_id_length: int = TOOL_CALL_ID_MAX_LEN,
+        put_inner_thoughts_in_kwargs: bool = False,
+        use_developer_message: bool = False,
+    ) -> dict:
+        """Go from Message class to ChatCompletion message object"""
+
+        # TODO change to pydantic casting, eg `return SystemMessageModel(self)`
+        # If we only have one content part and it's text, treat it as COT
+        parse_content_parts = False
+        if (
+            self.content
+            and len(self.content) == 1
+            and isinstance(self.content[0], TextContent)
+        ):
+            content = self.content[0].text
+        # Otherwise, check if we have TextContent and multiple other parts
+        elif self.content and len(self.content) > 1:
+            content = []
+            text_content_count = 0
+
+            for content_part in self.content:
+                if isinstance(content_part, TextContent):
+                    content.append(
+                        {
+                            "type": "text",
+                            "text": content_part.text,
+                        }
+                    )
+                    text_content_count += 1
+                elif isinstance(content_part, ImageContent):
+                    content.append(
+                        {
+                            "type": content_part.type,
+                            "image_id": content_part.image_id,
+                            "detail": content_part.detail,
+                        }
+                    )
+                elif isinstance(content_part, FileContent):
+                    content.append(
+                        {
+                            "type": content_part.type,
+                            "file_id": content_part.file_id,
+                        }
+                    )
+                elif isinstance(content_part, CloudFileContent):
+                    content.append(
+                        {
+                            "type": content_part.type,
+                            "cloud_file_uri": content_part.cloud_file_uri,
+                        }
+                    )
+                else:
+                    raise ValueError(f"Invalid content type: {content_part.type}")
+
+            if text_content_count > 1:
+                # TODO: (yu) @caren check this
+                parse_content_parts = True
+        else:
+            content = None
+
+        # TODO(caren) we should eventually support multiple content parts here?
+        # ie, actually make dict['content'] type list
+        # But for now, it's OK until we support multi-modal,
+        # since the only "parts" we have are for supporting various COT
+
+        if self.role == "system":
+            assert all([v is not None for v in [self.role]]), vars(self)
+            openai_message = {
+                "content": content,
+                "role": "developer" if use_developer_message else self.role,
+            }
+
+        elif self.role == "user":
+            assert all([v is not None for v in [content, self.role]]), vars(self)
+            openai_message = {
+                "content": content,
+                "role": self.role,
+            }
+
+        elif self.role == "assistant":
+            assert self.tool_calls is not None or content is not None
+            openai_message = {
+                "content": None if put_inner_thoughts_in_kwargs else content,
+                "role": self.role,
+            }
+
+            if self.tool_calls is not None:
+                openai_message["tool_calls"] = [
+                    tool_call.model_dump() for tool_call in self.tool_calls
+                ]
+                if max_tool_id_length:
+                    for tool_call_dict in openai_message["tool_calls"]:
+                        tool_call_dict["id"] = tool_call_dict["id"][:max_tool_id_length]
+
+        elif self.role == "tool":
+            assert all([v is not None for v in [self.role, self.tool_call_id]]), vars(
+                self
+            )
+            openai_message = {
+                "content": content,
+                "role": self.role,
+                "tool_call_id": self.tool_call_id[:max_tool_id_length]
+                if max_tool_id_length
+                else self.tool_call_id,
+            }
+
+        else:
+            raise ValueError(self.role)
+
+        if parse_content_parts:
+            for content in self.content:
+                if isinstance(content, ReasoningContent):
+                    openai_message["reasoning_content"] = content.reasoning
+                    if content.signature:
+                        openai_message["reasoning_content_signature"] = (
+                            content.signature
+                        )
+                if isinstance(content, RedactedReasoningContent):
+                    openai_message["redacted_reasoning_content"] = content.data
+
+        return openai_message
+
+    def to_anthropic_dict(
+        self,
+        inner_thoughts_xml_tag="thinking",
+        put_inner_thoughts_in_kwargs: bool = False,
+    ) -> dict:
+        """
+        Convert to an Anthropic message dictionary
+
+        Args:
+            inner_thoughts_xml_tag (str): The XML tag to wrap around inner thoughts
+        """
+
+        # Check for COT
+        if (
+            self.content
+            and len(self.content) == 1
+            and isinstance(self.content[0], TextContent)
+        ):
+            content = self.content[0].text
+        elif self.content and len(self.content) > 1:
+            assert self.role == "user"
+            content = []
+
+            for content_part in self.content:
+                if isinstance(content_part, TextContent):
+                    content.append(
+                        {
+                            "type": "text",
+                            "text": content_part.text,
+                        }
+                    )
+                elif isinstance(content_part, ImageContent):
+                    content.append(
+                        {"type": "image_url", "image_id": content_part.image_id}
+                    )
+                elif isinstance(content_part, FileContent):
+                    content.append(
+                        {
+                            "type": "file_uri",
+                            "file_id": content_part.file_id,
+                        }
+                    )
+                elif isinstance(content_part, CloudFileContent):
+                    content.append(
+                        {
+                            "type": "cloud_file_uri",
+                            "cloud_file_uri": content_part.cloud_file_uri,
+                        }
+                    )
+
+        else:
+            content = None
+
+        def add_xml_tag(string: str, xml_tag: Optional[str]):
+            # NOTE: Anthropic docs recommends using <thinking> tag when using CoT + tool use
+            if f"<{xml_tag}>" in string and f"</{xml_tag}>" in string:
+                # don't nest if tags already exist
+                return string
+            return f"<{xml_tag}>{string}</{xml_tag}" if xml_tag else string
+
+        if self.role == "system":
+            # NOTE: this is not for system instructions, but instead system "events"
+
+            assert all([v is not None for v in [content, self.role]]), vars(self)
+            # Two options here, we would use system.package_system_message,
+            # or use a more Anthropic-specific packaging ie xml tags
+            user_system_event = add_xml_tag(
+                string=f"SYSTEM ALERT: {content}", xml_tag="event"
+            )
+            anthropic_message = {
+                "content": user_system_event,
+                "role": "user",
+            }
+
+        elif self.role == "user":
+            assert all([v is not None for v in [content, self.role]]), vars(self)
+            anthropic_message = {
+                "content": content,
+                "role": self.role,
+            }
+
+        elif self.role == "assistant":
+            assert self.tool_calls is not None or content is not None
+            anthropic_message = {
+                "role": self.role,
+            }
+            content = []
+            # COT / reasoning / thinking
+            if len(self.content) > 1:
+                for content_part in self.content:
+                    if isinstance(content_part, ReasoningContent):
+                        content.append(
+                            {
+                                "type": "thinking",
+                                "thinking": content_part.reasoning,
+                                "signature": content_part.signature,
+                            }
+                        )
+                    if isinstance(content_part, RedactedReasoningContent):
+                        content.append(
+                            {
+                                "type": "redacted_thinking",
+                                "data": content_part.data,
+                            }
+                        )
+            elif content is not None:
+                content.append(
+                    {
+                        "type": "text",
+                        "text": add_xml_tag(
+                            string=content, xml_tag=inner_thoughts_xml_tag
+                        ),
+                    }
+                )
+            # Tool calling
+            if self.tool_calls is not None:
+                for tool_call in self.tool_calls:
+                    tool_call_input = parse_json(tool_call.function.arguments)
+
+                    content.append(
+                        {
+                            "type": "tool_use",
+                            "id": tool_call.id,
+                            "name": tool_call.function.name,
+                            "input": tool_call_input,
+                        }
+                    )
+
+            # If the only content was text, unpack it back into a singleton
+            # TODO support multi-modal
+            anthropic_message["content"] = content
+
+        elif self.role == "tool":
+            # NOTE: Anthropic uses role "user" for "tool" responses
+            assert all([v is not None for v in [self.role, self.tool_call_id]]), vars(
+                self
+            )
+            anthropic_message = {
+                "role": "user",  # NOTE: diff
+                "content": [
+                    # TODO support error types etc
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": self.tool_call_id,
+                        "content": content,
+                    }
+                ],
+            }
+
+        else:
+            raise ValueError(self.role)
+
+        return anthropic_message
+
+    def to_google_ai_dict(self, put_inner_thoughts_in_kwargs: bool = True) -> dict:
+        """
+        Go from Message class to Google AI REST message object
+        """
+        # type Content: https://ai.google.dev/api/rest/v1/Content / https://ai.google.dev/api/rest/v1beta/Content
+        #     parts[]: Part
+        #     role: str ('user' or 'model')
+        if (
+            self.content
+            and len(self.content) == 1
+            and isinstance(self.content[0], TextContent)
+        ):
+            text_content = self.content[0].text
+            contents = [{"text": text_content}]
+        elif self.content:
+            assert self.role == "user"
+            contents = []
+            for content in self.content:
+                if isinstance(content, ImageContent):
+                    contents.append({"image_id": content.image_id})
+                elif isinstance(content, TextContent):
+                    contents.append({"text": content.text})
+                elif isinstance(content, FileContent):
+                    contents.append({"file_id": content.file_id})
+                elif isinstance(content, CloudFileContent):
+                    contents.append({"cloud_file_uri": content.cloud_file_uri})
+                else:
+                    raise ValueError(f"Invalid content type: {content.type}")
+        else:
+            text_content = None
+            contents = None
+
+        if self.role != "tool" and self.name is not None:
+            warnings.warn(
+                f"Using Google AI with non-null 'name' field (name={self.name} role={self.role}), not yet supported."
+            )
+
+        if self.role == "system":
+            # NOTE: Gemini API doesn't have a 'system' role, use 'user' instead
+            # https://www.reddit.com/r/Bard/comments/1b90i8o/does_gemini_have_a_system_prompt_option_while/
+            google_ai_message = {
+                "role": "user",  # NOTE: no 'system'
+                "parts": [{"text": text_content}],
+            }
+
+        elif self.role == "user":
+            if not all([v is not None for v in [contents, self.role]]):
+                import ipdb
+
+                ipdb.set_trace()
+
+            assert all([v is not None for v in [contents, self.role]]), vars(self)
+            google_ai_message = {
+                "role": "user",
+                "parts": contents,
+            }
+
+        elif self.role == "assistant":
+            assert self.tool_calls is not None or text_content is not None
+            google_ai_message = {
+                "role": "model",  # NOTE: different
+            }
+
+            # NOTE: Google AI API doesn't allow non-null content + function call
+            # To get around this, just two a two part message, inner thoughts first then
+            parts = []
+            if not put_inner_thoughts_in_kwargs and text_content is not None:
+                # NOTE: ideally we do multi-part for CoT / inner thoughts + function call, but Google AI API doesn't allow it
+                raise NotImplementedError
+                parts.append({"text": text_content})
+
+            if self.tool_calls is not None:
+                # NOTE: implied support for multiple calls
+                for tool_call in self.tool_calls:
+                    function_name = tool_call.function.name
+                    function_args = tool_call.function.arguments
+                    try:
+                        # NOTE: Google AI wants actual JSON objects, not strings
+                        function_args = parse_json(function_args)
+                    except Exception:
+                        raise UserWarning(
+                            f"Failed to parse JSON function args: {function_args}"
+                        )
+                    
+                    parts.append(
+                        {
+                            "functionCall": {
+                                "name": function_name,
+                                "args": function_args,
+                            }
+                        }
+                    )
+            else:
+                assert text_content is not None
+                parts.append({"text": text_content})
+            google_ai_message["parts"] = parts
+
+        elif self.role == "tool":
+            # NOTE: Significantly different tool calling format, more similar to function calling format
+            assert all([v is not None for v in [self.role, self.tool_call_id]]), vars(
+                self
+            )
+
+            if self.name is None:
+                warnings.warn(
+                    "Couldn't find function name on tool call, defaulting to tool ID instead."
+                )
+                function_name = self.tool_call_id
+            else:
+                function_name = self.name
+
+            # NOTE: Google AI API wants the function response as JSON only, no string
+            try:
+                function_response = parse_json(text_content)
+            except Exception:
+                function_response = {"function_response": text_content}
+
+            google_ai_message = {
+                "role": "function",
+                "parts": [
+                    {
+                        "functionResponse": {
+                            "name": function_name,
+                            "response": {
+                                "name": function_name,  # NOTE: name twice... why?
+                                "content": function_response,
+                            },
+                        }
+                    }
+                ],
+            }
+
+        else:
+            raise ValueError(self.role)
+
+        # Validate that parts is never empty before returning
+        if "parts" not in google_ai_message or not google_ai_message["parts"]:
+            # If parts is empty, add a default text part
+            google_ai_message["parts"] = [{"text": "empty message"}]
+            warnings.warn(
+                f"Empty 'parts' detected in message with role '{self.role}'. Added default empty text part. Full message:\n{vars(self)}"
+            )
+
+        return google_ai_message
+
+    def to_cohere_dict(
+        self,
+        function_call_role: Optional[str] = "SYSTEM",
+        function_call_prefix: Optional[str] = "[CHATBOT called function]",
+        function_response_role: Optional[str] = "SYSTEM",
+        function_response_prefix: Optional[str] = "[CHATBOT function returned]",
+        inner_thoughts_as_kwarg: Optional[bool] = False,
+    ) -> List[dict]:
+        """
+        Cohere chat_history dicts only have 'role' and 'message' fields
+        """
+
+        # NOTE: returns a list of dicts so that we can convert:
+        #  assistant [cot]: "I'll send a message"
+        #  assistant [func]: send_message("hi")
+        #  tool: {'status': 'OK'}
+        # to:
+        #  CHATBOT.text: "I'll send a message"
+        #  SYSTEM.text: [CHATBOT called function] send_message("hi")
+        #  SYSTEM.text: [CHATBOT function returned] {'status': 'OK'}
+
+        # TODO: update this prompt style once guidance from Cohere on
+        # embedded function calls in multi-turn conversation become more clear
+        if (
+            self.content
+            and len(self.content) == 1
+            and isinstance(self.content[0], TextContent)
+        ):
+            text_content = self.content[0].text
+        else:
+            text_content = None
+        if self.role == "system":
+            """
+            The chat_history parameter should not be used for SYSTEM messages in most cases.
+            Instead, to add a SYSTEM role message at the beginning of a conversation, the preamble parameter should be used.
+            """
+            raise UserWarning(
+                "role 'system' messages should go in 'preamble' field for Cohere API"
+            )
+
+        elif self.role == "user":
+            assert all([v is not None for v in [text_content, self.role]]), vars(self)
+            cohere_message = [
+                {
+                    "role": "USER",
+                    "message": text_content,
+                }
+            ]
+
+        elif self.role == "assistant":
+            # NOTE: we may break this into two message - an inner thought and a function call
+            # Optionally, we could just make this a function call with the inner thought inside
+            assert self.tool_calls is not None or text_content is not None
+
+            if text_content and self.tool_calls:
+                if inner_thoughts_as_kwarg:
+                    raise NotImplementedError
+                cohere_message = [
+                    {
+                        "role": "CHATBOT",
+                        "message": text_content,
+                    },
+                ]
+                for tc in self.tool_calls:
+                    function_name = tc.function["name"]
+                    function_args = parse_json(tc.function["arguments"])
+                    function_args_str = ",".join(
+                        [f"{k}={v}" for k, v in function_args.items()]
+                    )
+                    function_call_text = f"{function_name}({function_args_str})"
+                    cohere_message.append(
+                        {
+                            "role": function_call_role,
+                            "message": f"{function_call_prefix} {function_call_text}",
+                        }
+                    )
+            elif not text_content and self.tool_calls:
+                cohere_message = []
+                for tc in self.tool_calls:
+                    # TODO better way to pack?
+                    function_call_text = json_dumps(tc.to_dict())
+                    cohere_message.append(
+                        {
+                            "role": function_call_role,
+                            "message": f"{function_call_prefix} {function_call_text}",
+                        }
+                    )
+            elif text_content and not self.tool_calls:
+                cohere_message = [
+                    {
+                        "role": "CHATBOT",
+                        "message": text_content,
+                    }
+                ]
+            else:
+                raise ValueError("Message does not have content nor tool_calls")
+
+        elif self.role == "tool":
+            assert all([v is not None for v in [self.role, self.tool_call_id]]), vars(
+                self
+            )
+            function_response_text = text_content
+            cohere_message = [
+                {
+                    "role": function_response_role,
+                    "message": f"{function_response_prefix} {function_response_text}",
+                }
+            ]
+
+        else:
+            raise ValueError(self.role)
+
+        return cohere_message
+
+    @staticmethod
+    def generate_otid_from_id(message_id: str, index: int) -> str:
+        """
+        Convert message id to bits and change the list bit to the index
+        """
+        if not 0 <= index < 128:
+            raise ValueError("Index must be between 0 and 127")
+
+        message_uuid = message_id.replace("message-", "")
+        uuid_int = int(message_uuid.replace("-", ""), 16)
+
+        # Clear last 7 bits and set them to index; supports up to 128 unique indices
+        uuid_int = (uuid_int & ~0x7F) | (index & 0x7F)
+
+        hex_str = f"{uuid_int:032x}"
+        return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:]}"
+
+
+class ToolReturn(BaseModel):
+    status: Literal["success", "error"] = Field(
+        ..., description="The status of the tool call"
+    )
+    stdout: Optional[List[str]] = Field(
+        None, description="Captured stdout (e.g. prints, logs) from the tool invocation"
+    )
+    stderr: Optional[List[str]] = Field(
+        None, description="Captured stderr from the tool invocation"
+    )