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

mirix/helpers/tool_rule_solver.py

@@ -0,0 +1,166 @@
+import json
+from typing import List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from mirix.schemas.enums import ToolRuleType
+from mirix.schemas.tool_rule import (
+    BaseToolRule,
+    ChildToolRule,
+    ConditionalToolRule,
+    InitToolRule,
+    TerminalToolRule,
+)
+
+
+class ToolRuleValidationError(Exception):
+    """Custom exception for tool rule validation errors in ToolRulesSolver."""
+
+    def __init__(self, message: str):
+        super().__init__(f"ToolRuleValidationError: {message}")
+
+
+class ToolRulesSolver(BaseModel):
+    init_tool_rules: List[InitToolRule] = Field(
+        default_factory=list,
+        description="Initial tool rules to be used at the start of tool execution.",
+    )
+    tool_rules: List[Union[ChildToolRule, ConditionalToolRule]] = Field(
+        default_factory=list,
+        description="Standard tool rules for controlling execution sequence and allowed transitions.",
+    )
+    terminal_tool_rules: List[TerminalToolRule] = Field(
+        default_factory=list,
+        description="Terminal tool rules that end the agent loop if called.",
+    )
+    last_tool_name: Optional[str] = Field(
+        None, description="The most recent tool used, updated with each tool call."
+    )
+
+    def __init__(self, tool_rules: List[BaseToolRule], **kwargs):
+        super().__init__(**kwargs)
+        # Separate the provided tool rules into init, standard, and terminal categories
+        for rule in tool_rules:
+            if rule.type == ToolRuleType.run_first:
+                assert isinstance(rule, InitToolRule)
+                self.init_tool_rules.append(rule)
+            elif rule.type == ToolRuleType.constrain_child_tools:
+                assert isinstance(rule, ChildToolRule)
+                self.tool_rules.append(rule)
+            elif rule.type == ToolRuleType.conditional:
+                assert isinstance(rule, ConditionalToolRule)
+                self.validate_conditional_tool(rule)
+                self.tool_rules.append(rule)
+            elif rule.type == ToolRuleType.exit_loop:
+                assert isinstance(rule, TerminalToolRule)
+                self.terminal_tool_rules.append(rule)
+
+    def update_tool_usage(self, tool_name: str):
+        """Update the internal state to track the last tool called."""
+        self.last_tool_name = tool_name
+
+    def get_allowed_tool_names(
+        self, error_on_empty: bool = False, last_function_response: Optional[str] = None
+    ) -> List[str]:
+        """Get a list of tool names allowed based on the last tool called."""
+        if self.last_tool_name is None:
+            # Use initial tool rules if no tool has been called yet
+            return [rule.tool_name for rule in self.init_tool_rules]
+        else:
+            # Find a matching ToolRule for the last tool used
+            current_rule = next(
+                (
+                    rule
+                    for rule in self.tool_rules
+                    if rule.tool_name == self.last_tool_name
+                ),
+                None,
+            )
+
+            if current_rule is None:
+                if error_on_empty:
+                    raise ValueError(f"No tool rule found for {self.last_tool_name}")
+                return []
+
+            # If the current rule is a conditional tool rule, use the LLM response to
+            # determine which child tool to use
+            if isinstance(current_rule, ConditionalToolRule):
+                if not last_function_response:
+                    raise ValueError(
+                        "Conditional tool rule requires an LLM response to determine which child tool to use"
+                    )
+                next_tool = self.evaluate_conditional_tool(
+                    current_rule, last_function_response
+                )
+                return [next_tool] if next_tool else []
+
+            return current_rule.children if current_rule.children else []
+
+    def is_terminal_tool(self, tool_name: str) -> bool:
+        """Check if the tool is defined as a terminal tool in the terminal tool rules."""
+        return any(rule.tool_name == tool_name for rule in self.terminal_tool_rules)
+
+    def has_children_tools(self, tool_name):
+        """Check if the tool has children tools"""
+        return any(rule.tool_name == tool_name for rule in self.tool_rules)
+
+    def validate_conditional_tool(self, rule: ConditionalToolRule):
+        """
+        Validate a conditional tool rule
+
+        Args:
+            rule (ConditionalToolRule): The conditional tool rule to validate
+
+        Raises:
+            ToolRuleValidationError: If the rule is invalid
+        """
+        if len(rule.child_output_mapping) == 0:
+            raise ToolRuleValidationError(
+                "Conditional tool rule must have at least one child tool."
+            )
+        return True
+
+    def evaluate_conditional_tool(
+        self, tool: ConditionalToolRule, last_function_response: str
+    ) -> str:
+        """
+        Parse function response to determine which child tool to use based on the mapping
+
+        Args:
+            tool (ConditionalToolRule): The conditional tool rule
+            last_function_response (str): The function response in JSON format
+
+        Returns:
+            str: The name of the child tool to use next
+        """
+        json_response = json.loads(last_function_response)
+        function_output = json_response["message"]
+
+        # Try to match the function output with a mapping key
+        for key in tool.child_output_mapping:
+            # Convert function output to match key type for comparison
+            if isinstance(key, bool):
+                typed_output = function_output.lower() == "true"
+            elif isinstance(key, int):
+                try:
+                    typed_output = int(function_output)
+                except (ValueError, TypeError):
+                    continue
+            elif isinstance(key, float):
+                try:
+                    typed_output = float(function_output)
+                except (ValueError, TypeError):
+                    continue
+            else:  # string
+                if function_output == "True" or function_output == "False":
+                    typed_output = function_output.lower()
+                elif function_output == "None":
+                    typed_output = None
+                else:
+                    typed_output = function_output
+
+            if typed_output == key:
+                return tool.child_output_mapping[key]
+
+        # If no match found, use default
+        return tool.default_child

mirix/schemas/__init__.py

@@ -0,0 +1 @@
+# Mirix schemas package

mirix/schemas/agent.py

@@ -0,0 +1,400 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator
+
+from mirix.constants import DEFAULT_EMBEDDING_CHUNK_SIZE
+from mirix.helpers import ToolRulesSolver
+from mirix.schemas.block import CreateBlock
+from mirix.schemas.embedding_config import EmbeddingConfig
+from mirix.schemas.llm_config import LLMConfig
+from mirix.schemas.memory import Memory
+from mirix.schemas.message import Message, MessageCreate
+from mirix.schemas.mirix_base import OrmMetadataBase
+from mirix.schemas.openai.chat_completion_response import UsageStatistics
+from mirix.schemas.tool import Tool
+from mirix.schemas.tool_rule import ToolRule
+from mirix.utils import create_random_username
+
+
+class AgentType(str, Enum):
+    """
+    Enum to represent the type of agent.
+    """
+
+    coder_agent = "coder_agent"
+    chat_agent = "chat_agent"
+    reflexion_agent = "reflexion_agent"
+    background_agent = "background_agent"
+    episodic_memory_agent = "episodic_memory_agent"
+    procedural_memory_agent = "procedural_memory_agent"
+    resource_memory_agent = "resource_memory_agent"
+    knowledge_vault_memory_agent = "knowledge_vault_memory_agent"
+    meta_memory_agent = "meta_memory_agent"
+    semantic_memory_agent = "semantic_memory_agent"
+    core_memory_agent = "core_memory_agent"
+
+
+class AgentState(OrmMetadataBase, validate_assignment=True):
+    """
+    Representation of an agent's state. This is the state of the agent at a given time, and is persisted in the DB backend. The state has all the information needed to recreate a persisted agent.
+
+    Parameters:
+        id (str): The unique identifier of the agent.
+        name (str): The name of the agent (must be unique to the user).
+        created_at (datetime): The datetime the agent was created.
+        memory (Memory): The in-context memory of the agent.
+        tools (List[str]): The tools used by the agent. This includes any memory editing functions specified in `memory`.
+        system (str): The system prompt used by the agent.
+        llm_config (LLMConfig): The LLM configuration used by the agent.
+        embedding_config (EmbeddingConfig): The embedding configuration used by the agent.
+
+    """
+
+    __id_prefix__ = "agent"
+
+    # NOTE: this is what is returned to the client and also what is used to initialize `Agent`
+    id: str = Field(..., description="The id of the agent. Assigned by the database.")
+    name: str = Field(..., description="The name of the agent.")
+    # tool rules
+    tool_rules: Optional[List[ToolRule]] = Field(
+        default=None, description="The list of tool rules."
+    )
+
+    # in-context memory
+    message_ids: Optional[List[str]] = Field(
+        default=None,
+        description="The ids of the messages in the agent's in-context memory.",
+    )
+
+    # system prompt
+    system: str = Field(..., description="The system prompt used by the agent.")
+
+    # agent configuration
+    agent_type: AgentType = Field(..., description="The type of agent.")
+
+    # llm information
+    llm_config: LLMConfig = Field(
+        ..., description="The LLM configuration used by the agent."
+    )
+    embedding_config: EmbeddingConfig = Field(
+        ..., description="The embedding configuration used by the agent."
+    )
+
+    # This is an object representing the in-process state of a running `Agent`
+    # Field in this object can be theoretically edited by tools, and will be persisted by the ORM
+    organization_id: Optional[str] = Field(
+        None,
+        description="The unique identifier of the organization associated with the agent.",
+    )
+
+    description: Optional[str] = Field(
+        None, description="The description of the agent."
+    )
+    parent_id: Optional[str] = Field(
+        None, description="The parent agent ID (for sub-agents in a meta-agent)."
+    )
+    children: Optional[List["AgentState"]] = Field(
+        default=None, description="Child agents (sub-agents) if this is a parent agent."
+    )
+
+    memory: Memory = Field(..., description="The in-context memory of the agent.")
+    tools: List[Tool] = Field(..., description="The tools used by the agent.")
+    mcp_tools: Optional[List[str]] = Field(
+        default_factory=list,
+        description="List of connected MCP server names (e.g., ['gmail-native'])",
+    )
+
+
+class CreateAgent(BaseModel, validate_assignment=True):  #
+    # all optional as server can generate defaults
+    name: str = Field(
+        default_factory=lambda: create_random_username(),
+        description="The name of the agent.",
+    )
+
+    # memory creation
+    memory_blocks: Optional[List[CreateBlock]] = Field(
+        None,
+        description="The blocks to create in the agent's in-context memory.",
+    )
+    # TODO: This is a legacy field and should be removed ASAP to force `tool_ids` usage
+    tools: Optional[List[str]] = Field(None, description="The tools used by the agent.")
+    tool_ids: Optional[List[str]] = Field(
+        None, description="The ids of the tools used by the agent."
+    )
+    tool_rules: Optional[List[ToolRule]] = Field(
+        None, description="The tool rules governing the agent."
+    )
+    system: Optional[str] = Field(
+        None, description="The system prompt used by the agent."
+    )
+    agent_type: AgentType = Field(
+        default_factory=lambda: AgentType.chat_agent, description="The type of agent."
+    )
+    llm_config: Optional[LLMConfig] = Field(
+        None, description="The LLM configuration used by the agent."
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the agent."
+    )
+    # Note: if this is None, then we'll populate with the standard "more human than human" initial message sequence
+    # If the client wants to make this empty, then the client can set the arg to an empty list
+    initial_message_sequence: Optional[List[MessageCreate]] = Field(
+        None,
+        description="The initial set of messages to put in the agent's in-context memory.",
+    )
+    include_base_tools: bool = Field(
+        True,
+        description="If true, attaches the Mirix core tools (e.g. archival_memory and core_memory related functions).",
+    )
+    include_multi_agent_tools: bool = Field(
+        False,
+        description="If true, attaches the Mirix multi-agent tools (e.g. sending a message to another agent).",
+    )
+    parent_id: Optional[str] = Field(
+        None, description="The parent agent ID (for sub-agents in a meta-agent)."
+    )
+    model: Optional[str] = Field(
+        None,
+        description="The LLM configuration handle used by the agent, specified in the format "
+        "provider/model-name, as an alternative to specifying llm_config.",
+    )
+    embedding: Optional[str] = Field(
+        None,
+        description="The embedding configuration handle used by the agent, specified in the format provider/model-name.",
+    )
+    context_window_limit: Optional[int] = Field(
+        None, description="The context window limit used by the agent."
+    )
+    embedding_chunk_size: Optional[int] = Field(
+        DEFAULT_EMBEDDING_CHUNK_SIZE,
+        description="The embedding chunk size used by the agent.",
+    )
+    from_template: Optional[str] = Field(
+        None, description="The template id used to configure the agent"
+    )
+    template: bool = Field(False, description="Whether the agent is a template")
+    project: Optional[str] = Field(
+        None, description="The project slug that the agent will be associated with."
+    )
+    tool_exec_environment_variables: Optional[Dict[str, str]] = Field(
+        None,
+        description="The environment variables for tool execution specific to this agent.",
+    )
+    memory_variables: Optional[Dict[str, str]] = Field(
+        None, description="The variables that should be set for the agent."
+    )
+    mcp_tools: Optional[List[str]] = Field(
+        None, description="List of MCP server names to connect to this agent."
+    )
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, name: str) -> str:
+        """Validate the requested new agent name (prevent bad inputs)"""
+
+        import re
+
+        if not name:
+            # don't check if not provided
+            return name
+
+        # TODO: this check should also be added to other model (e.g. User.name)
+        # Length check
+        if not (1 <= len(name) <= 50):
+            raise ValueError("Name length must be between 1 and 50 characters.")
+
+        # Regex for allowed characters (alphanumeric, spaces, hyphens, underscores)
+        if not re.match("^[A-Za-z0-9 _-]+$", name):
+            raise ValueError("Name contains invalid characters.")
+
+        # Further checks can be added here...
+        # TODO
+
+        return name
+
+    @field_validator("model")
+    @classmethod
+    def validate_model(cls, model: Optional[str]) -> Optional[str]:
+        if not model:
+            return model
+
+        provider_name, model_name = model.split("/", 1)
+        if not provider_name or not model_name:
+            raise ValueError(
+                "The llm config handle should be in the format provider/model-name"
+            )
+
+        return model
+
+    @field_validator("embedding")
+    @classmethod
+    def validate_embedding(cls, embedding: Optional[str]) -> Optional[str]:
+        if not embedding:
+            return embedding
+
+        provider_name, embedding_name = embedding.split("/", 1)
+        if not provider_name or not embedding_name:
+            raise ValueError(
+                "The embedding config handle should be in the format provider/model-name"
+            )
+
+        return embedding
+
+
+class UpdateAgent(BaseModel):
+    name: Optional[str] = Field(None, description="The name of the agent.")
+    tool_ids: Optional[List[str]] = Field(
+        None, description="The ids of the tools used by the agent."
+    )
+    block_ids: Optional[List[str]] = Field(
+        None, description="The ids of the blocks used by the agent."
+    )
+    system: Optional[str] = Field(
+        None, description="The system prompt used by the agent."
+    )
+    tool_rules: Optional[List[ToolRule]] = Field(
+        None, description="The tool rules governing the agent."
+    )
+    llm_config: Optional[LLMConfig] = Field(
+        None, description="The LLM configuration used by the agent."
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None, description="The embedding configuration used by the agent."
+    )
+    message_ids: Optional[List[str]] = Field(
+        None, description="The ids of the messages in the agent's in-context memory."
+    )
+    description: Optional[str] = Field(
+        None, description="The description of the agent."
+    )
+    parent_id: Optional[str] = Field(
+        None, description="The parent agent ID (for sub-agents in a meta-agent)."
+    )
+    mcp_tools: Optional[List[str]] = Field(
+        None, description="List of MCP server names to connect to this agent."
+    )
+
+    class Config:
+        extra = "ignore"  # Ignores extra fields
+
+
+class CreateMetaAgent(BaseModel):
+    """Request schema for creating a MetaAgent."""
+
+    name: Optional[str] = Field(
+        None,
+        description="Optional name for the MetaAgent. If None, a random name will be generated.",
+    )
+    agents: List[Union[str, Dict[str, Any]]] = Field(
+        default_factory=lambda: [
+            "core_memory_agent",
+            "resource_memory_agent",
+            "semantic_memory_agent",
+            "episodic_memory_agent",
+            "procedural_memory_agent",
+            "knowledge_vault_memory_agent",
+            "meta_memory_agent",
+            "reflexion_agent",
+            "background_agent",
+        ],
+        description="List of memory agent names or dicts with agent configs. Supports both 'agent_name' strings and {'agent_name': {'blocks': [...], ...}} dicts.",
+    )
+    system_prompts: Optional[Dict[str, str]] = Field(
+        None,
+        description="Dictionary mapping agent names to their system prompt text. Takes precedence over system_prompts_folder.",
+    )
+    llm_config: Optional[LLMConfig] = Field(
+        None,
+        description="LLM configuration for memory agents. Required if no default is set.",
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None,
+        description="Embedding configuration for memory agents. Required if no default is set.",
+    )
+
+class UpdateMetaAgent(BaseModel):
+    """Request schema for updating a MetaAgent."""
+
+    name: Optional[str] = Field(
+        None,
+        description="Optional new name for the MetaAgent.",
+    )
+    agents: Optional[List[Union[str, Dict[str, Any]]]] = Field(
+        None,
+        description="List of memory agent names or dicts with agent configs. Will be compared with existing agents to determine what to add/remove.",
+    )
+    system_prompts: Optional[Dict[str, str]] = Field(
+        None,
+        description="Dictionary mapping agent names to their system prompt text. Updates only the specified agents.",
+    )
+    llm_config: Optional[LLMConfig] = Field(
+        None,
+        description="LLM configuration for meta agent and its sub-agents.",
+    )
+    embedding_config: Optional[EmbeddingConfig] = Field(
+        None,
+        description="Embedding configuration for meta agent and its sub-agents.",
+    )
+
+    class Config:
+        extra = "ignore"  # Ignores extra fields
+
+
+class AgentStepResponse(BaseModel):
+    messages: List[Message] = Field(
+        ..., description="The messages generated during the agent's step."
+    )
+    continue_chaining: bool = Field(
+        ...,
+        description="Whether the agent requested a contine_chaining (i.e. follow-up execution).",
+    )
+    function_failed: bool = Field(
+        ..., description="Whether the agent step ended because a function call failed."
+    )
+    in_context_memory_warning: bool = Field(
+        ...,
+        description="Whether the agent step ended because the in-context memory is near its limit.",
+    )
+    usage: UsageStatistics = Field(
+        ..., description="Usage statistics of the LLM call during the agent's step."
+    )
+    traj: Optional[dict] = Field(
+        None, description="Action, Observation, State at the current step"
+    )
+
+
+class AgentStepState(BaseModel):
+    step_number: int = Field(
+        ..., description="The current step number in the agent loop"
+    )
+    tool_rules_solver: ToolRulesSolver = Field(
+        ..., description="The current state of the ToolRulesSolver"
+    )
+
+
+def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None):
+    if agent_type == AgentType.sleeptime_agent:
+        return (
+            "{% for block in blocks %}"
+            '<{{ block.label }} characters="{{ block.value|length }}/{{ block.limit }}">\n'
+            "{% for line in block.value.split('\\n') %}"
+            "Line {{ loop.index }}: {{ line }}\n"
+            "{% endfor %}"
+            "</{{ block.label }}>"
+            "{% if not loop.last %}\n{% endif %}"
+            "{% endfor %}"
+        )
+    return (
+        "{% for block in blocks %}"
+        '<{{ block.label }} characters="{{ block.value|length }}/{{ block.limit }}">\n'
+        "{{ block.value }}\n"
+        "</{{ block.label }}>"
+        "{% if not loop.last %}\n{% endif %}"
+        "{% endfor %}"
+    )
+
+
+# Rebuild model to support forward references (children field)
+AgentState.model_rebuild()