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

mirix/schemas/mirix_message_content.py

@@ -0,0 +1,230 @@
+from enum import Enum
+from typing import Annotated, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class MessageContentType(str, Enum):
+    text = "text"
+    image_url = "image_url"
+    file_uri = "file_uri"
+    google_cloud_file_uri = "google_cloud_file_uri"
+    tool_call = "tool_call"
+    tool_return = "tool_return"
+    reasoning = "reasoning"
+    redacted_reasoning = "redacted_reasoning"
+    omitted_reasoning = "omitted_reasoning"
+
+
+class MessageContent(BaseModel):
+    type: MessageContentType = Field(..., description="The type of the message.")
+
+
+# -------------------------------
+# User Content Types
+# -------------------------------
+
+
+class TextContent(MessageContent):
+    type: Literal[MessageContentType.text] = Field(
+        MessageContentType.text, description="The type of the message."
+    )
+    text: str = Field(..., description="The text content of the message.")
+
+
+class ImageContent(MessageContent):
+    type: Literal[MessageContentType.image_url] = Field(
+        MessageContentType.image_url, description="The type of the message."
+    )
+    image_id: str = Field(..., description="The id of the image in the database")
+    detail: Optional[str] = Field(None, description="The detail of the image.")
+
+
+class FileContent(MessageContent):
+    type: Literal[MessageContentType.file_uri] = Field(
+        MessageContentType.file_uri, description="The type of the message."
+    )
+    file_id: str = Field(..., description="The id of the file in the database")
+
+
+class CloudFileContent(MessageContent):
+    type: Literal[MessageContentType.google_cloud_file_uri] = Field(
+        MessageContentType.google_cloud_file_uri, description="The type of the message."
+    )
+    cloud_file_uri: str = Field(..., description="The URI of the file in the database")
+
+
+MirixUserMessageContentUnion = Annotated[
+    Union[TextContent, ImageContent, FileContent, CloudFileContent],
+    Field(discriminator="type"),
+]
+
+
+def create_mirix_user_message_content_union_schema():
+    return {
+        "oneOf": [
+            {"$ref": "#/components/schemas/TextContent"},
+        ],
+        "discriminator": {
+            "propertyName": "type",
+            "mapping": {
+                "text": "#/components/schemas/TextContent",
+                "image": "#/components/schemas/ImageContent",  # TODO: (yu) Not sure about this
+            },
+        },
+    }
+
+
+def get_mirix_user_message_content_union_str_json_schema():
+    return {
+        "anyOf": [
+            {
+                "type": "array",
+                "items": {
+                    "$ref": "#/components/schemas/MirixUserMessageContentUnion",
+                },
+            },
+            {"type": "string"},
+        ],
+    }
+
+
+# -------------------------------
+# Assistant Content Types
+# -------------------------------
+
+
+MirixAssistantMessageContentUnion = Annotated[
+    Union[TextContent],
+    Field(discriminator="type"),
+]
+
+
+def create_mirix_assistant_message_content_union_schema():
+    return {
+        "oneOf": [
+            {"$ref": "#/components/schemas/TextContent"},
+        ],
+        "discriminator": {
+            "propertyName": "type",
+            "mapping": {
+                "text": "#/components/schemas/TextContent",
+            },
+        },
+    }
+
+
+def get_mirix_assistant_message_content_union_str_json_schema():
+    return {
+        "anyOf": [
+            {
+                "type": "array",
+                "items": {
+                    "$ref": "#/components/schemas/MirixAssistantMessageContentUnion",
+                },
+            },
+            {"type": "string"},
+        ],
+    }
+
+
+# -------------------------------
+# Intermediate Step Content Types
+# -------------------------------
+
+
+class ToolCallContent(MessageContent):
+    type: Literal[MessageContentType.tool_call] = Field(
+        MessageContentType.tool_call,
+        description="Indicates this content represents a tool call event.",
+    )
+    id: str = Field(
+        ..., description="A unique identifier for this specific tool call instance."
+    )
+    name: str = Field(..., description="The name of the tool being called.")
+    input: dict = Field(
+        ...,
+        description="The parameters being passed to the tool, structured as a dictionary of parameter names to values.",
+    )
+
+
+class ToolReturnContent(MessageContent):
+    type: Literal[MessageContentType.tool_return] = Field(
+        MessageContentType.tool_return,
+        description="Indicates this content represents a tool return event.",
+    )
+    tool_call_id: str = Field(
+        ...,
+        description="References the ID of the ToolCallContent that initiated this tool call.",
+    )
+    content: str = Field(..., description="The content returned by the tool execution.")
+    is_error: bool = Field(
+        ..., description="Indicates whether the tool execution resulted in an error."
+    )
+
+
+class ReasoningContent(MessageContent):
+    type: Literal[MessageContentType.reasoning] = Field(
+        MessageContentType.reasoning,
+        description="Indicates this is a reasoning/intermediate step.",
+    )
+    is_native: bool = Field(
+        ...,
+        description="Whether the reasoning content was generated by a reasoner model that processed this step.",
+    )
+    reasoning: str = Field(
+        ..., description="The intermediate reasoning or thought process content."
+    )
+    signature: Optional[str] = Field(
+        None, description="A unique identifier for this reasoning step."
+    )
+
+
+class RedactedReasoningContent(MessageContent):
+    type: Literal[MessageContentType.redacted_reasoning] = Field(
+        MessageContentType.redacted_reasoning,
+        description="Indicates this is a redacted thinking step.",
+    )
+    data: str = Field(
+        ..., description="The redacted or filtered intermediate reasoning content."
+    )
+
+
+class OmittedReasoningContent(MessageContent):
+    type: Literal[MessageContentType.omitted_reasoning] = Field(
+        MessageContentType.omitted_reasoning,
+        description="Indicates this is an omitted reasoning step.",
+    )
+    tokens: int = Field(
+        ..., description="The reasoning token count for intermediate reasoning content."
+    )
+
+
+MirixMessageContentUnion = Annotated[
+    Union[
+        TextContent,
+        ImageContent,
+        FileContent,
+        CloudFileContent,
+        ToolCallContent,
+        ToolReturnContent,
+        ReasoningContent,
+        RedactedReasoningContent,
+        OmittedReasoningContent,
+    ],
+    Field(discriminator="type"),
+]
+
+
+def get_mirix_message_content_union_str_json_schema():
+    return {
+        "anyOf": [
+            {
+                "type": "array",
+                "items": {
+                    "$ref": "#/components/schemas/MiriMessageContentUnion",
+                },
+            },
+            {"type": "string"},
+        ],
+    }

mirix/schemas/mirix_request.py

@@ -0,0 +1,39 @@
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from mirix.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
+from mirix.schemas.message import MessageCreate
+
+
+class MirixRequestConfig(BaseModel):
+    # Flags to support the use of AssistantMessage message types
+    use_assistant_message: bool = Field(
+        default=True,
+        description="Whether the server should parse specific tool call arguments (default `send_message`) as `AssistantMessage` objects.",
+    )
+    assistant_message_tool_name: str = Field(
+        default=DEFAULT_MESSAGE_TOOL,
+        description="The name of the designated message tool.",
+    )
+    assistant_message_tool_kwarg: str = Field(
+        default=DEFAULT_MESSAGE_TOOL_KWARG,
+        description="The name of the message argument in the designated message tool.",
+    )
+
+
+class MirixRequest(BaseModel):
+    messages: List[MessageCreate] = Field(
+        ..., description="The messages to be sent to the agent."
+    )
+    config: MirixRequestConfig = Field(
+        default=MirixRequestConfig(),
+        description="Configuration options for the MirixRequest.",
+    )
+
+
+class MirixStreamingRequest(MirixRequest):
+    stream_tokens: bool = Field(
+        default=False,
+        description="Flag to determine if individual tokens should be streamed. Set to True for token streaming (requires stream_steps = True).",
+    )

mirix/schemas/mirix_response.py

@@ -0,0 +1,183 @@
+import html
+import json
+import re
+from typing import List, Union
+
+from pydantic import BaseModel, Field
+
+from mirix.schemas.enums import MessageStreamStatus
+from mirix.schemas.mirix_message import MirixMessage, MirixMessageUnion
+from mirix.schemas.usage import MirixUsageStatistics
+from mirix.utils import json_dumps
+
+# TODO: consider moving into own file
+
+
+class MirixResponse(BaseModel):
+    """
+    Response object from an agent interaction, consisting of the new messages generated by the agent and usage statistics.
+    The type of the returned messages can be either `Message` or `MirixMessage`, depending on what was specified in the request.
+
+    Attributes:
+        messages (List[Union[Message, MirixMessage]]): The messages returned by the agent.
+        usage (MirixUsageStatistics): The usage statistics
+    """
+
+    messages: List[MirixMessageUnion] = Field(
+        ...,
+        description="The messages returned by the agent.",
+        json_schema_extra={
+            "items": {
+                "$ref": "#/components/schemas/MirixMessageUnion",
+            }
+        },
+    )
+    usage: MirixUsageStatistics = Field(
+        ...,
+        description="The usage statistics of the agent.",
+    )
+
+    def __str__(self):
+        return json_dumps(
+            {
+                "messages": [message.model_dump() for message in self.messages],
+                # Assume `Message` and `MirixMessage` have a `dict()` method
+                "usage": self.usage.model_dump(),  # Assume `MirixUsageStatistics` has a `dict()` method
+            },
+            indent=4,
+        )
+
+    def _repr_html_(self):
+        def get_formatted_content(msg):
+            if msg.message_type == "internal_monologue":
+                return f'<div class="content"><span class="internal-monologue">{html.escape(msg.internal_monologue)}</span></div>'
+            if msg.message_type == "reasoning_message":
+                return f'<div class="content"><span class="internal-monologue">{html.escape(msg.reasoning)}</span></div>'
+            elif msg.message_type == "function_call":
+                args = format_json(msg.function_call.arguments)
+                return f'<div class="content"><span class="function-name">{html.escape(msg.function_call.name)}</span>({args})</div>'
+            elif msg.message_type == "tool_call_message":
+                args = format_json(msg.tool_call.arguments)
+                return f'<div class="content"><span class="function-name">{html.escape(msg.tool_call.name)}</span>({args})</div>'
+            elif msg.message_type == "function_return":
+                return_value = format_json(msg.function_return)
+                # return f'<div class="status-line">Status: {html.escape(msg.status)}</div><div class="content">{return_value}</div>'
+                return f'<div class="content">{return_value}</div>'
+            elif msg.message_type == "tool_return_message":
+                return_value = format_json(msg.tool_return)
+                # return f'<div class="status-line">Status: {html.escape(msg.status)}</div><div class="content">{return_value}</div>'
+                return f'<div class="content">{return_value}</div>'
+            elif msg.message_type == "user_message":
+                if is_json(msg.message):
+                    return f'<div class="content">{format_json(msg.message)}</div>'
+                else:
+                    return f'<div class="content">{html.escape(msg.message)}</div>'
+            elif msg.message_type in ["assistant_message", "system_message"]:
+                return f'<div class="content">{html.escape(msg.message)}</div>'
+            else:
+                return f'<div class="content">{html.escape(str(msg))}</div>'
+
+        def is_json(string):
+            try:
+                json.loads(string)
+                return True
+            except ValueError:
+                return False
+
+        def format_json(json_str):
+            try:
+                parsed = json.loads(json_str)
+                formatted = json.dumps(parsed, indent=2, ensure_ascii=False)
+                formatted = (
+                    formatted.replace("&", "&amp;")
+                    .replace("<", "&lt;")
+                    .replace(">", "&gt;")
+                )
+                formatted = formatted.replace("\n", "<br>").replace(
+                    "  ", "&nbsp;&nbsp;"
+                )
+                formatted = re.sub(
+                    r'(".*?"):', r'<span class="json-key">\1</span>:', formatted
+                )
+                formatted = re.sub(
+                    r': (".*?")', r': <span class="json-string">\1</span>', formatted
+                )
+                formatted = re.sub(
+                    r": (\d+)", r': <span class="json-number">\1</span>', formatted
+                )
+                formatted = re.sub(
+                    r": (true|false)",
+                    r': <span class="json-boolean">\1</span>',
+                    formatted,
+                )
+                return formatted
+            except json.JSONDecodeError:
+                return html.escape(json_str)
+
+        html_output = """
+        <style>
+            .message-container, .usage-container {
+                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+                max-width: 800px;
+                margin: 20px auto;
+                background-color: #1e1e1e;
+                border-radius: 8px;
+                overflow: hidden;
+                color: #d4d4d4;
+            }
+            .message, .usage-stats {
+                padding: 10px 15px;
+                border-bottom: 1px solid #3a3a3a;
+            }
+            .message:last-child, .usage-stats:last-child {
+                border-bottom: none;
+            }
+            .title {
+                font-weight: bold;
+                margin-bottom: 5px;
+                color: #ffffff;
+                text-transform: uppercase;
+                font-size: 0.9em;
+            }
+            .content {
+                background-color: #2d2d2d;
+                border-radius: 4px;
+                padding: 5px 10px;
+                font-family: 'Consolas', 'Courier New', monospace;
+                white-space: pre-wrap;
+            }
+            .json-key, .function-name, .json-boolean { color: #9cdcfe; }
+            .json-string { color: #ce9178; }
+            .json-number { color: #b5cea8; }
+            .internal-monologue { font-style: italic; }
+        </style>
+        <div class="message-container">
+        """
+
+        for msg in self.messages:
+            content = get_formatted_content(msg)
+            title = msg.message_type.replace("_", " ").upper()
+            html_output += f"""
+            <div class="message">
+                <div class="title">{title}</div>
+                {content}
+            </div>
+            """
+        html_output += "</div>"
+
+        # Formatting the usage statistics
+        usage_html = json.dumps(self.usage.model_dump(), indent=2)
+        html_output += f"""
+        <div class="usage-container">
+            <div class="usage-stats">
+                <div class="title">USAGE STATISTICS</div>
+                <div class="content">{format_json(usage_html)}</div>
+            </div>
+        </div>
+        """
+
+        return html_output
+
+
+# The streaming response is either [DONE], [DONE_STEP], [DONE], an error, or a MirixMessage
+MirixStreamingResponse = Union[MirixMessage, MessageStreamStatus, MirixUsageStatistics]

mirix/schemas/openai/__init__.py

@@ -0,0 +1 @@
+# OpenAI schema package

mirix/schemas/openai/chat_completion_request.py

@@ -0,0 +1,122 @@
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class SystemMessage(BaseModel):
+    content: str
+    role: str = "system"
+    name: Optional[str] = None
+
+
+class UserMessage(BaseModel):
+    content: Union[str, List[dict]]
+    role: str = "user"
+    name: Optional[str] = None
+
+
+class ToolCallFunction(BaseModel):
+    name: str
+    arguments: str
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: Literal["function"] = "function"
+    function: ToolCallFunction
+
+
+class AssistantMessage(BaseModel):
+    content: Optional[str] = None
+    role: str = "assistant"
+    name: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class ToolMessage(BaseModel):
+    content: str
+    role: str = "tool"
+    tool_call_id: str
+
+
+ChatMessage = Union[SystemMessage, UserMessage, AssistantMessage, ToolMessage]
+
+
+# TODO: this might not be necessary with the validator
+def cast_message_to_subtype(m_dict: dict) -> ChatMessage:
+    """Cast a dictionary to one of the individual message types"""
+    role = m_dict.get("role")
+    if role == "system" or role == "developer":
+        return SystemMessage(**m_dict)
+    elif role == "user":
+        return UserMessage(**m_dict)
+    elif role == "assistant":
+        return AssistantMessage(**m_dict)
+    elif role == "tool":
+        return ToolMessage(**m_dict)
+    else:
+        raise ValueError(f"Unknown message role: {role}")
+
+
+class ResponseFormat(BaseModel):
+    type: str = Field(default="text", pattern="^(text|json_object)$")
+
+
+## tool_choice ##
+class FunctionCall(BaseModel):
+    name: str
+
+
+class ToolFunctionChoice(BaseModel):
+    # The type of the tool. Currently, only function is supported
+    type: Literal["function"] = "function"
+    # type: str = Field(default="function", const=True)
+    function: FunctionCall
+
+
+ToolChoice = Union[Literal["none", "auto", "required"], ToolFunctionChoice]
+
+
+## tools ##
+class FunctionSchema(BaseModel):
+    name: str
+    description: Optional[str] = None
+    parameters: Optional[Dict[str, Any]] = None  # JSON Schema for the parameters
+
+
+class Tool(BaseModel):
+    # The type of the tool. Currently, only function is supported
+    type: Literal["function"] = "function"
+    # type: str = Field(default="function", const=True)
+    function: FunctionSchema
+
+
+## function_call ##
+FunctionCallChoice = Union[Literal["none", "auto"], FunctionCall]
+
+
+class ChatCompletionRequest(BaseModel):
+    """https://platform.openai.com/docs/api-reference/chat/create"""
+
+    model: str
+    messages: List[ChatMessage]
+    frequency_penalty: Optional[float] = 0
+    logit_bias: Optional[Dict[str, int]] = None
+    logprobs: Optional[bool] = False
+    top_logprobs: Optional[int] = None
+    max_tokens: Optional[int] = None
+    n: Optional[int] = 1
+    presence_penalty: Optional[float] = 0
+    response_format: Optional[ResponseFormat] = None
+    seed: Optional[int] = None
+    stop: Optional[Union[str, List[str]]] = None
+    stream: Optional[bool] = False
+    temperature: Optional[float] = 1  # TODO: might need to add logics to control this
+    top_p: Optional[float] = 1
+
+    # function-calling related
+    tools: Optional[List[Tool]] = None
+    tool_choice: Optional[ToolChoice] = None  # "none" means don't call a tool
+    # deprecated scheme
+    functions: Optional[List[FunctionSchema]] = None
+    function_call: Optional[FunctionCallChoice] = None