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

mirix/schemas/openai/chat_completion_response.py

@@ -0,0 +1,144 @@
+import datetime
+from typing import Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel
+
+# class ToolCallFunction(BaseModel):
+#     name: str
+#     arguments: str
+
+
+class FunctionCall(BaseModel):
+    arguments: str
+    name: str
+
+
+class ToolCall(BaseModel):
+    id: str
+    # "Currently, only function is supported"
+    type: Literal["function"] = "function"
+    # function: ToolCallFunction
+    function: FunctionCall
+
+
+class LogProbToken(BaseModel):
+    token: str
+    logprob: float
+    bytes: Optional[List[int]]
+
+
+class MessageContentLogProb(BaseModel):
+    token: str
+    logprob: float
+    bytes: Optional[List[int]]
+    top_logprobs: Optional[List[LogProbToken]]
+
+
+class Message(BaseModel):
+    content: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+    role: str
+    function_call: Optional[FunctionCall] = None  # Deprecated
+
+
+class Choice(BaseModel):
+    finish_reason: str
+    index: int
+    message: Message
+    logprobs: Optional[Dict[str, Union[List[MessageContentLogProb], None]]] = None
+    seed: Optional[int] = None  # found in TogetherAI
+
+
+class UsageStatistics(BaseModel):
+    completion_tokens: int = 0
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    last_prompt_tokens: int = 0
+    last_completion_tokens: int = 0
+
+    def __add__(self, other: "UsageStatistics") -> "UsageStatistics":
+        return UsageStatistics(
+            completion_tokens=self.completion_tokens + other.completion_tokens,
+            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+            total_tokens=self.total_tokens + other.total_tokens,
+        )
+
+
+class ChatCompletionResponse(BaseModel):
+    """https://platform.openai.com/docs/api-reference/chat/object"""
+
+    id: str
+    choices: List[Choice]
+    created: datetime.datetime
+    model: Optional[str] = (
+        None  # NOTE: this is not consistent with OpenAI API standard, however is necessary to support local LLMs
+    )
+    # system_fingerprint: str  # docs say this is mandatory, but in reality API returns None
+    system_fingerprint: Optional[str] = None
+    # object: str = Field(default="chat.completion")
+    object: Literal["chat.completion"] = "chat.completion"
+    usage: UsageStatistics
+
+    def __str__(self):
+        return self.model_dump_json(indent=4)
+
+
+class FunctionCallDelta(BaseModel):
+    # arguments: Optional[str] = None
+    name: Optional[str] = None
+    arguments: str
+    # name: str
+
+
+class ToolCallDelta(BaseModel):
+    index: int
+    id: Optional[str] = None
+    # "Currently, only function is supported"
+    type: Literal["function"] = "function"
+    # function: ToolCallFunction
+    function: Optional[FunctionCallDelta] = None
+
+
+class MessageDelta(BaseModel):
+    """Partial delta stream of a Message
+
+    Example ChunkResponse:
+    {
+        'id': 'chatcmpl-9EOCkKdicNo1tiL1956kPvCnL2lLS',
+        'object': 'chat.completion.chunk',
+        'created': 1713216662,
+        'model': 'gpt-4-0613',
+        'system_fingerprint': None,
+        'choices': [{
+            'index': 0,
+            'delta': {'content': 'User'},
+            'logprobs': None,
+            'finish_reason': None
+        }]
+    }
+    """
+
+    content: Optional[str] = None
+    tool_calls: Optional[List[ToolCallDelta]] = None
+    # role: Optional[str] = None
+    function_call: Optional[FunctionCallDelta] = None  # Deprecated
+
+
+class ChunkChoice(BaseModel):
+    finish_reason: Optional[str] = None  # NOTE: when streaming will be null
+    index: int
+    delta: MessageDelta
+    logprobs: Optional[Dict[str, Union[List[MessageContentLogProb], None]]] = None
+
+
+class ChatCompletionChunkResponse(BaseModel):
+    """https://platform.openai.com/docs/api-reference/chat/streaming"""
+
+    id: str
+    choices: List[ChunkChoice]
+    created: datetime.datetime
+    model: str
+    # system_fingerprint: str  # docs say this is mandatory, but in reality API returns None
+    system_fingerprint: Optional[str] = None
+    # object: str = Field(default="chat.completion")
+    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"

mirix/schemas/openai/chat_completions.py

@@ -0,0 +1,127 @@
+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[str]]
+    role: str = "user"
+    name: Optional[str] = None
+
+
+class ToolCallFunction(BaseModel):
+    name: str = Field(..., description="The name of the function to call")
+    arguments: str = Field(
+        ..., description="The arguments to pass to the function (JSON dump)"
+    )
+
+
+class ToolCall(BaseModel):
+    id: str = Field(..., description="The ID of the tool call")
+    type: str = "function"
+    function: ToolCallFunction = Field(
+        ..., description="The arguments and name for the function"
+    )
+
+
+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":
+        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("Unknown message 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"], 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
+    top_p: Optional[float] = 1
+    user: Optional[str] = None  # unique ID of the end-user (for monitoring)
+
+    # function-calling related
+    tools: Optional[List[Tool]] = None
+    tool_choice: Optional[ToolChoice] = "none"
+    # deprecated scheme
+    functions: Optional[List[FunctionSchema]] = None
+    function_call: Optional[FunctionCallChoice] = None

mirix/schemas/openai/embedding_response.py

@@ -0,0 +1,11 @@
+from typing import List, Literal
+
+from pydantic import BaseModel
+
+
+class EmbeddingResponse(BaseModel):
+    """OpenAI embedding response model: https://platform.openai.com/docs/api-reference/embeddings/object"""
+
+    index: int  # the index of the embedding in the list of embeddings
+    embedding: List[float]
+    object: Literal["embedding"] = "embedding"

mirix/schemas/openai/openai.py

@@ -0,0 +1,229 @@
+from enum import Enum
+from typing import Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class ImageFile(BaseModel):
+    type: str = "image_file"
+    file_id: str
+
+
+class Text(BaseModel):
+    object: str = "text"
+    text: str = Field(..., description="The text content to be processed by the agent.")
+
+
+class MessageRoleType(str, Enum):
+    user = "user"
+    system = "system"
+
+
+class OpenAIAssistant(BaseModel):
+    """Represents an OpenAI assistant (equivalent to Mirix preset)"""
+
+    id: str = Field(..., description="The unique identifier of the assistant.")
+    name: str = Field(..., description="The name of the assistant.")
+    object: str = "assistant"
+    description: Optional[str] = Field(
+        None, description="The description of the assistant."
+    )
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the assistant was created."
+    )
+    model: str = Field(..., description="The model used by the assistant.")
+    instructions: str = Field(..., description="The instructions for the assistant.")
+    tools: Optional[List[str]] = Field(
+        None, description="The tools used by the assistant."
+    )
+    file_ids: Optional[List[str]] = Field(
+        None, description="List of file IDs associated with the assistant."
+    )
+    metadata: Optional[dict] = Field(
+        None, description="Metadata associated with the assistant."
+    )
+
+
+class OpenAIMessage(BaseModel):
+    id: str = Field(..., description="The unique identifier of the message.")
+    object: str = "thread.message"
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the message was created."
+    )
+    thread_id: str = Field(..., description="The unique identifier of the thread.")
+    role: str = Field(
+        ..., description="Role of the message sender (either 'user' or 'system')"
+    )
+    content: List[Union[Text, ImageFile]] = Field(
+        None, description="The message content to be processed by the agent."
+    )
+    assistant_id: str = Field(
+        ..., description="The unique identifier of the assistant."
+    )
+    run_id: Optional[str] = Field(None, description="The unique identifier of the run.")
+    file_ids: Optional[List[str]] = Field(
+        None, description="List of file IDs associated with the message."
+    )
+    metadata: Optional[Dict] = Field(
+        None, description="Metadata associated with the message."
+    )
+
+
+class OpenAIThread(BaseModel):
+    """Represents an OpenAI thread (equivalent to Mirix agent)"""
+
+    id: str = Field(..., description="The unique identifier of the thread.")
+    object: str = "thread"
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the thread was created."
+    )
+    metadata: dict = Field(None, description="Metadata associated with the thread.")
+
+
+class AssistantFile(BaseModel):
+    id: str = Field(..., description="The unique identifier of the file.")
+    object: str = "assistant.file"
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the file was created."
+    )
+    assistant_id: str = Field(
+        ..., description="The unique identifier of the assistant."
+    )
+
+
+class MessageFile(BaseModel):
+    id: str = Field(..., description="The unique identifier of the file.")
+    object: str = "thread.message.file"
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the file was created."
+    )
+    message_id: str = Field(..., description="The unique identifier of the message.")
+
+
+class Function(BaseModel):
+    name: str = Field(..., description="The name of the function.")
+    arguments: str = Field(..., description="The arguments of the function.")
+
+
+class ToolCall(BaseModel):
+    id: str = Field(..., description="The unique identifier of the tool call.")
+    type: str = "function"
+    function: Function = Field(..., description="The function call.")
+
+
+class ToolCallOutput(BaseModel):
+    tool_call_id: str = Field(
+        ..., description="The unique identifier of the tool call."
+    )
+    output: str = Field(..., description="The output of the tool call.")
+
+
+class RequiredAction(BaseModel):
+    type: str = "submit_tool_outputs"
+    submit_tool_outputs: List[ToolCall]
+
+
+class OpenAIError(BaseModel):
+    code: str = Field(..., description="The error code.")
+    message: str = Field(..., description="The error message.")
+
+
+class OpenAIUsage(BaseModel):
+    completion_tokens: int = Field(
+        ..., description="The number of tokens used for the run."
+    )
+    prompt_tokens: int = Field(
+        ..., description="The number of tokens used for the prompt."
+    )
+    total_tokens: int = Field(
+        ..., description="The total number of tokens used for the run."
+    )
+
+
+class OpenAIMessageCreationStep(BaseModel):
+    type: str = "message_creation"
+    message_id: str = Field(..., description="The unique identifier of the message.")
+
+
+class OpenAIToolCallsStep(BaseModel):
+    type: str = "tool_calls"
+    tool_calls: List[ToolCall] = Field(..., description="The tool calls.")
+
+
+class OpenAIRun(BaseModel):
+    id: str = Field(..., description="The unique identifier of the run.")
+    object: str = "thread.run"
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the run was created."
+    )
+    thread_id: str = Field(..., description="The unique identifier of the thread.")
+    assistant_id: str = Field(
+        ..., description="The unique identifier of the assistant."
+    )
+    status: str = Field(..., description="The status of the run.")
+    required_action: Optional[RequiredAction] = Field(
+        None, description="The required action of the run."
+    )
+    last_error: Optional[OpenAIError] = Field(
+        None, description="The last error of the run."
+    )
+    expires_at: int = Field(
+        ..., description="The unix timestamp of when the run expires."
+    )
+    started_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run started."
+    )
+    cancelled_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run was cancelled."
+    )
+    failed_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run failed."
+    )
+    completed_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run completed."
+    )
+    model: str = Field(..., description="The model used by the run.")
+    instructions: str = Field(..., description="The instructions for the run.")
+    tools: Optional[List[ToolCall]] = Field(
+        None, description="The tools used by the run."
+    )  # TODO: also add code interpreter / retrieval
+    file_ids: Optional[List[str]] = Field(
+        None, description="List of file IDs associated with the run."
+    )
+    metadata: Optional[dict] = Field(
+        None, description="Metadata associated with the run."
+    )
+    usage: Optional[OpenAIUsage] = Field(None, description="The usage of the run.")
+
+
+class OpenAIRunStep(BaseModel):
+    id: str = Field(..., description="The unique identifier of the run step.")
+    object: str = "thread.run.step"
+    created_at: int = Field(
+        ..., description="The unix timestamp of when the run step was created."
+    )
+    assistant_id: str = Field(
+        ..., description="The unique identifier of the assistant."
+    )
+    thread_id: str = Field(..., description="The unique identifier of the thread.")
+    run_id: str = Field(..., description="The unique identifier of the run.")
+    type: str = Field(
+        ..., description="The type of the run step."
+    )  # message_creation, tool_calls
+    status: str = Field(..., description="The status of the run step.")
+    step_defaults: Union[OpenAIToolCallsStep, OpenAIMessageCreationStep] = Field(
+        ..., description="The step defaults."
+    )
+    last_error: Optional[OpenAIError] = Field(
+        None, description="The last error of the run step."
+    )
+    expired_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run step expired."
+    )
+    failed_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run failed."
+    )
+    completed_at: Optional[int] = Field(
+        None, description="The unix timestamp of when the run completed."
+    )
+    usage: Optional[OpenAIUsage] = Field(None, description="The usage of the run.")

mirix/schemas/organization.py

@@ -0,0 +1,38 @@
+from datetime import datetime
+from typing import Optional
+import uuid
+
+from pydantic import Field
+
+from mirix.schemas.mirix_base import MirixBase
+from mirix.utils import create_random_username, get_utc_time
+
+
+class OrganizationBase(MirixBase):
+    __id_prefix__ = "org"
+
+
+def _generate_org_id() -> str:
+    """Generate a random organization ID."""
+    return f"org-{uuid.uuid4().hex[:8]}"
+
+
+class Organization(OrganizationBase):
+    id: str = Field(
+        default_factory=_generate_org_id,
+        description="The unique identifier of the organization.",
+    )
+    name: str = Field(
+        create_random_username(),
+        description="The name of the organization.",
+        json_schema_extra={"default": "SincereYogurt"},
+    )
+    created_at: Optional[datetime] = Field(
+        default_factory=get_utc_time,
+        description="The creation date of the organization.",
+    )
+
+
+class OrganizationCreate(OrganizationBase):
+    id: Optional[str] = Field(None, description="The unique identifier of the organization.")
+    name: Optional[str] = Field(None, description="The name of the organization.")

mirix/schemas/procedural_memory.py

@@ -0,0 +1,151 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field, field_validator
+
+from mirix.constants import MAX_EMBEDDING_DIM
+from mirix.schemas.embedding_config import EmbeddingConfig
+from mirix.schemas.mirix_base import MirixBase
+from mirix.utils import get_utc_time
+
+
+class ProceduralMemoryItemBase(MirixBase):
+    """
+    Base schema for storing procedural knowledge (e.g., workflows, methods).
+    """
+
+    __id_prefix__ = "proc_item"
+    entry_type: str = Field(
+        ..., description="Category (e.g., 'workflow', 'guide', 'script')"
+    )
+    summary: str = Field(..., description="Short descriptive text about the procedure")
+    steps: List[str] = Field(
+        ..., description="Step-by-step instructions as a list of strings"
+    )
+
+
+class ProceduralMemoryItem(ProceduralMemoryItemBase):
+    """
+    Full procedural memory item schema, with database-related fields.
+    """
+
+    id: Optional[str] = Field(
+        None, description="Unique identifier for the procedural memory item"
+    )
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this procedural memory item belongs to"
+    )
+    user_id: str = Field(
+        ..., description="The id of the user who generated the procedure"
+    )
+    created_at: datetime = Field(
+        default_factory=get_utc_time, description="Creation timestamp"
+    )
+    updated_at: Optional[datetime] = Field(None, description="Last update timestamp")
+    last_modify: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "timestamp": get_utc_time().isoformat(),
+            "operation": "created",
+        },
+        description="Last modification info including timestamp and operation type",
+    )
+    organization_id: str = Field(
+        ..., description="The unique identifier of the organization"
+    )
+    summary_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the summary"
+    )
+    steps_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the steps"
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the event"
+    )
+    
+    # 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"
+            }
+        ]
+    )
+
+    # need to validate both steps_embedding and summary_embedding to ensure they are the same size    
+    # 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("summary_embedding", "steps_embedding")
+    @classmethod
+    def pad_embeddings(cls, embedding: List[float]) -> List[float]:
+        """Pad embeddings to `MAX_EMBEDDING_SIZE`. This is necessary to ensure all stored embeddings are the same size."""
+        import numpy as np
+
+        if embedding and len(embedding) != MAX_EMBEDDING_DIM:
+            np_embedding = np.array(embedding)
+            padded_embedding = np.pad(
+                np_embedding,
+                (0, MAX_EMBEDDING_DIM - np_embedding.shape[0]),
+                mode="constant",
+            )
+            return padded_embedding.tolist()
+        return embedding
+
+
+class ProceduralMemoryItemUpdate(MirixBase):
+    """Schema for updating an existing procedural memory item."""
+
+    id: str = Field(..., description="Unique ID for this procedural memory entry")
+    agent_id: Optional[str] = Field(
+        None, description="The id of the agent this procedural memory item belongs to"
+    )
+    entry_type: Optional[str] = Field(
+        None, description="Category (e.g., 'workflow', 'guide', 'script')"
+    )
+    summary: Optional[str] = Field(None, description="Short descriptive text")
+    steps: Optional[List[str]] = Field(
+        None, description="Step-by-step instructions as a list of strings"
+    )
+    organization_id: Optional[str] = Field(None, description="The organization ID")
+    updated_at: datetime = Field(
+        default_factory=get_utc_time, description="Update timestamp"
+    )
+    last_modify: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Last modification info including timestamp and operation type",
+    )
+    steps_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the event"
+    )
+    summary_embedding: Optional[List[float]] = Field(
+        None, description="The embedding of the summary"
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the event"
+    )
+
+
+    filter_tags: Optional[Dict[str, Any]] = Field(
+        None, description="Custom filter tags for filtering and categorization"
+    )
+
+class ProceduralMemoryItemResponse(ProceduralMemoryItem):
+    """Response schema for procedural memory item."""
+
+    pass