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

mirix/schemas/tool.py

@@ -0,0 +1,241 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from pydantic import Field, model_validator
+
+from mirix.constants import (
+    COMPOSIO_TOOL_TAG_NAME,
+    FUNCTION_RETURN_CHAR_LIMIT,
+    MIRIX_CORE_TOOL_MODULE_NAME,
+    MIRIX_EXTRA_TOOL_MODULE_NAME,
+    MIRIX_MEMORY_TOOL_MODULE_NAME,
+)
+from mirix.functions.functions import (
+    derive_openai_json_schema,
+    get_json_schema_from_module,
+)
+from mirix.functions.helpers import generate_langchain_tool_wrapper
+from mirix.functions.schema_generator import generate_schema_from_args_schema_v2
+from mirix.orm.enums import ToolType
+from mirix.schemas.mirix_base import MirixBase
+
+if TYPE_CHECKING:
+    try:
+        from langchain_core.tools import BaseTool as LangChainBaseTool
+    except ImportError:
+        LangChainBaseTool = Any  # type: ignore
+
+
+class BaseTool(MirixBase):
+    __id_prefix__ = "tool"
+
+
+class Tool(BaseTool):
+    """
+    Representation of a tool, which is a function that can be called by the agent.
+
+    Parameters:
+        id (str): The unique identifier of the tool.
+        name (str): The name of the function.
+        tags (List[str]): Metadata tags.
+        source_code (str): The source code of the function.
+        json_schema (Dict): The JSON schema of the function.
+
+    """
+
+    id: str = BaseTool.generate_id_field()
+    tool_type: ToolType = Field(ToolType.CUSTOM, description="The type of the tool.")
+    description: Optional[str] = Field(None, description="The description of the tool.")
+    source_type: Optional[str] = Field(None, description="The type of the source code.")
+    organization_id: Optional[str] = Field(
+        None,
+        description="The unique identifier of the organization associated with the tool.",
+    )
+    name: Optional[str] = Field(None, description="The name of the function.")
+    tags: List[str] = Field([], description="Metadata tags.")
+
+    # code
+    source_code: Optional[str] = Field(
+        None, description="The source code of the function."
+    )
+    json_schema: Optional[Dict] = Field(
+        None, description="The JSON schema of the function."
+    )
+
+    # tool configuration
+    return_char_limit: int = Field(
+        FUNCTION_RETURN_CHAR_LIMIT,
+        description="The maximum number of characters in the response.",
+    )
+
+    # metadata fields
+    created_by_id: Optional[str] = Field(
+        None, description="The id of the user that made this Tool."
+    )
+    last_updated_by_id: Optional[str] = Field(
+        None, description="The id of the user that made this Tool."
+    )
+
+    @model_validator(mode="after")
+    def populate_missing_fields(self):
+        """
+        Populate missing fields: name, description, and json_schema.
+        """
+
+        if self.tool_type == ToolType.CUSTOM:
+            # If it's a custom tool, we need to ensure source_code is present
+            if not self.source_code:
+                raise ValueError(
+                    f"Custom tool with id={self.id} is missing source_code field."
+                )
+
+            # Always derive json_schema for freshest possible json_schema
+            # TODO: Instead of checking the tag, we should having `COMPOSIO` as a specific ToolType
+            # TODO: We skip this for Composio bc composio json schemas are derived differently
+            if COMPOSIO_TOOL_TAG_NAME not in self.tags:
+                self.json_schema = derive_openai_json_schema(
+                    source_code=self.source_code
+                )
+        elif self.tool_type in {ToolType.MIRIX_CORE}:
+            # If it's mirix core tool, we generate the json_schema on the fly here
+            self.json_schema = get_json_schema_from_module(
+                module_name=MIRIX_CORE_TOOL_MODULE_NAME, function_name=self.name
+            )
+        elif self.tool_type in {ToolType.MIRIX_MEMORY_CORE}:
+            self.json_schema = get_json_schema_from_module(
+                module_name=MIRIX_MEMORY_TOOL_MODULE_NAME, function_name=self.name
+            )
+        elif self.tool_type in {ToolType.MIRIX_EXTRA}:
+            self.json_schema = get_json_schema_from_module(
+                module_name=MIRIX_EXTRA_TOOL_MODULE_NAME, function_name=self.name
+            )
+        elif self.tool_type in {ToolType.MIRIX_MCP}:
+            # MCP tools have their json_schema already provided by MCP tool registry
+            # Skip validation since these are auto-generated tools
+            if not self.json_schema:
+                raise ValueError(f"MCP tool {self.name} is missing json_schema field")
+
+        # Derive name from the JSON schema if not provided
+        if not self.name:
+            # TODO: This in theory could error, but name should always be on json_schema
+            # TODO: Make JSON schema a typed pydantic object
+            self.name = self.json_schema.get("name")
+
+        # Derive description from the JSON schema if not provided
+        if not self.description:
+            # TODO: This in theory could error, but description should always be on json_schema
+            # TODO: Make JSON schema a typed pydantic object
+            self.description = self.json_schema.get("description")
+
+        return self
+
+
+class ToolCreate(MirixBase):
+    name: Optional[str] = Field(
+        None,
+        description="The name of the function (auto-generated from source_code if not provided).",
+    )
+    description: Optional[str] = Field(None, description="The description of the tool.")
+    tags: List[str] = Field([], description="Metadata tags.")
+    source_code: str = Field(..., description="The source code of the function.")
+    source_type: str = Field("python", description="The source type of the function.")
+    json_schema: Optional[Dict] = Field(
+        None,
+        description="The JSON schema of the function (auto-generated from source_code if not provided)",
+    )
+    return_char_limit: int = Field(
+        FUNCTION_RETURN_CHAR_LIMIT,
+        description="The maximum number of characters in the response.",
+    )
+
+    @classmethod
+    def from_langchain(
+        cls,
+        langchain_tool: "LangChainBaseTool",
+        additional_imports_module_attr_map: dict[str, str] = None,
+    ) -> "ToolCreate":
+        """
+        Class method to create an instance of Tool from a Langchain tool (must be from langchain_community.tools).
+
+        Args:
+            langchain_tool (LangChainBaseTool): An instance of a LangChain BaseTool (BaseTool from LangChain)
+            additional_imports_module_attr_map (dict[str, str]): A mapping of module names to attribute name. This is used internally to import all the required classes for the langchain tool. For example, you would pass in `{"langchain_community.utilities": "WikipediaAPIWrapper"}` for `from langchain_community.tools import WikipediaQueryRun`. NOTE: You do NOT need to specify the tool import here, that is done automatically for you.
+
+        Returns:
+            Tool: A Mirix Tool initialized with attributes derived from the provided LangChain BaseTool object.
+        """
+        description = langchain_tool.description
+        source_type = "python"
+        tags = ["langchain"]
+        # NOTE: langchain tools may come from different packages
+        wrapper_func_name, wrapper_function_str = generate_langchain_tool_wrapper(
+            langchain_tool, additional_imports_module_attr_map
+        )
+        json_schema = generate_schema_from_args_schema_v2(
+            langchain_tool.args_schema, name=wrapper_func_name, description=description
+        )
+
+        return cls(
+            name=wrapper_func_name,
+            description=description,
+            source_type=source_type,
+            tags=tags,
+            source_code=wrapper_function_str,
+            json_schema=json_schema,
+        )
+
+    @classmethod
+    def load_default_langchain_tools(cls) -> List["ToolCreate"]:
+        # For now, we only support wikipedia tool
+        from langchain_community.tools import WikipediaQueryRun
+        from langchain_community.utilities import WikipediaAPIWrapper
+
+        wikipedia_tool = ToolCreate.from_langchain(
+            WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()),
+            {"langchain_community.utilities": "WikipediaAPIWrapper"},
+        )
+
+        return [wikipedia_tool]
+
+    @classmethod
+    def load_default_composio_tools(cls) -> List["ToolCreate"]:
+        pass
+
+        # TODO: Disable composio tools for now
+        # TODO: Naming is causing issues
+        # calculator = ToolCreate.from_composio(action_name=Action.MATHEMATICAL_CALCULATOR.name)
+        # serp_news = ToolCreate.from_composio(action_name=Action.SERPAPI_NEWS_SEARCH.name)
+        # serp_google_search = ToolCreate.from_composio(action_name=Action.SERPAPI_SEARCH.name)
+        # serp_google_maps = ToolCreate.from_composio(action_name=Action.SERPAPI_GOOGLE_MAPS_SEARCH.name)
+
+        return []
+
+
+class ToolUpdate(MirixBase):
+    description: Optional[str] = Field(None, description="The description of the tool.")
+    name: Optional[str] = Field(None, description="The name of the function.")
+    tags: Optional[List[str]] = Field(None, description="Metadata tags.")
+    source_code: Optional[str] = Field(
+        None, description="The source code of the function."
+    )
+    source_type: Optional[str] = Field(None, description="The type of the source code.")
+    json_schema: Optional[Dict] = Field(
+        None,
+        description="The JSON schema of the function (auto-generated from source_code if not provided)",
+    )
+    return_char_limit: Optional[int] = Field(
+        None, description="The maximum number of characters in the response."
+    )
+
+    class Config:
+        extra = "ignore"  # Allows extra fields without validation errors
+        # TODO: Remove this, and clean usage of ToolUpdate everywhere else
+
+
+class ToolRunFromSource(MirixBase):
+    source_code: str = Field(..., description="The source code of the function.")
+    args: Dict[str, Any] = Field(..., description="The arguments to pass to the tool.")
+    env_vars: Dict[str, str] = Field(
+        None, description="The environment variables to pass to the tool."
+    )
+    name: Optional[str] = Field(None, description="The name of the tool to run.")
+    source_type: Optional[str] = Field(None, description="The type of the source code.")

mirix/schemas/tool_rule.py

@@ -0,0 +1,209 @@
+import json
+from typing import Annotated, Any, Dict, List, Literal, Optional, Set, Union
+
+from pydantic import Field
+
+from mirix.schemas.enums import ToolRuleType
+from mirix.schemas.mirix_base import MirixBase
+
+
+class BaseToolRule(MirixBase):
+    __id_prefix__ = "tool_rule"
+    tool_name: str = Field(
+        ...,
+        description="The name of the tool. Must exist in the database for the user's organization.",
+    )
+    type: ToolRuleType = Field(..., description="The type of the message.")
+
+    def get_valid_tools(
+        self,
+        tool_call_history: List[str],
+        available_tools: Set[str],
+        last_function_response: Optional[str],
+    ) -> set[str]:
+        raise NotImplementedError
+
+
+class ChildToolRule(BaseToolRule):
+    """
+    A ToolRule represents a tool that can be invoked by the agent.
+    """
+
+    type: Literal[ToolRuleType.constrain_child_tools] = (
+        ToolRuleType.constrain_child_tools
+    )
+    children: List[str] = Field(
+        ..., description="The children tools that can be invoked."
+    )
+
+    def get_valid_tools(
+        self,
+        tool_call_history: List[str],
+        available_tools: Set[str],
+        last_function_response: Optional[str],
+    ) -> Set[str]:
+        last_tool = tool_call_history[-1] if tool_call_history else None
+        return set(self.children) if last_tool == self.tool_name else available_tools
+
+
+class ParentToolRule(BaseToolRule):
+    """
+    A ToolRule that only allows a child tool to be called if the parent has been called.
+    """
+
+    type: Literal[ToolRuleType.parent_last_tool] = ToolRuleType.parent_last_tool
+    children: List[str] = Field(
+        ..., description="The children tools that can be invoked."
+    )
+
+    def get_valid_tools(
+        self,
+        tool_call_history: List[str],
+        available_tools: Set[str],
+        last_function_response: Optional[str],
+    ) -> Set[str]:
+        last_tool = tool_call_history[-1] if tool_call_history else None
+        return (
+            set(self.children)
+            if last_tool == self.tool_name
+            else available_tools - set(self.children)
+        )
+
+
+class ConditionalToolRule(BaseToolRule):
+    """
+    A ToolRule that conditionally maps to different child tools based on the output.
+    """
+
+    type: Literal[ToolRuleType.conditional] = ToolRuleType.conditional
+    default_child: Optional[str] = Field(
+        None,
+        description="The default child tool to be called. If None, any tool can be called.",
+    )
+    child_output_mapping: Dict[Any, str] = Field(
+        ..., description="The output case to check for mapping"
+    )
+    require_output_mapping: bool = Field(
+        default=False,
+        description="Whether to throw an error when output doesn't match any case",
+    )
+
+    def get_valid_tools(
+        self,
+        tool_call_history: List[str],
+        available_tools: Set[str],
+        last_function_response: Optional[str],
+    ) -> Set[str]:
+        """Determine valid tools based on function output mapping."""
+        if not tool_call_history or tool_call_history[-1] != self.tool_name:
+            return available_tools  # No constraints if this rule doesn't apply
+
+        if not last_function_response:
+            raise ValueError(
+                "Conditional tool rule requires an LLM response to determine which child tool to use"
+            )
+
+        try:
+            json_response = json.loads(last_function_response)
+            function_output = json_response.get("message", "")
+        except json.JSONDecodeError:
+            if self.require_output_mapping:
+                return set()  # Strict mode: Invalid response means no allowed tools
+            return {self.default_child} if self.default_child else available_tools
+
+        # Match function output to a mapped child tool
+        for key, tool in self.child_output_mapping.items():
+            if self._matches_key(function_output, key):
+                return {tool}
+
+        # If no match found, use default or allow all tools if no default is set
+        if self.require_output_mapping:
+            return set()  # Strict mode: No match means no valid tools
+
+        return {self.default_child} if self.default_child else available_tools
+
+    def _matches_key(self, function_output: str, key: Any) -> bool:
+        """Helper function to determine if function output matches a mapping key."""
+        if isinstance(key, bool):
+            return (
+                function_output.lower() == "true"
+                if key
+                else function_output.lower() == "false"
+            )
+        elif isinstance(key, int):
+            try:
+                return int(function_output) == key
+            except ValueError:
+                return False
+        elif isinstance(key, float):
+            try:
+                return float(function_output) == key
+            except ValueError:
+                return False
+        else:  # Assume string
+            return str(function_output) == str(key)
+
+
+class InitToolRule(BaseToolRule):
+    """
+    Represents the initial tool rule configuration.
+    """
+
+    type: Literal[ToolRuleType.run_first] = ToolRuleType.run_first
+
+
+class TerminalToolRule(BaseToolRule):
+    """
+    Represents a terminal tool rule configuration where if this tool gets called, it must end the agent loop.
+    """
+
+    type: Literal[ToolRuleType.exit_loop] = ToolRuleType.exit_loop
+
+
+class ContinueToolRule(BaseToolRule):
+    """
+    Represents a tool rule configuration where if this tool gets called, it must continue the agent loop.
+    """
+
+    type: Literal[ToolRuleType.continue_loop] = ToolRuleType.continue_loop
+
+
+class MaxCountPerStepToolRule(BaseToolRule):
+    """
+    Represents a tool rule configuration which constrains the total number of times this tool can be invoked in a single step.
+    """
+
+    type: Literal[ToolRuleType.max_count_per_step] = ToolRuleType.max_count_per_step
+    max_count_limit: int = Field(
+        ...,
+        description="The max limit for the total number of times this tool can be invoked in a single step.",
+    )
+
+    def get_valid_tools(
+        self,
+        tool_call_history: List[str],
+        available_tools: Set[str],
+        last_function_response: Optional[str],
+    ) -> Set[str]:
+        """Restricts the tool if it has been called max_count_limit times in the current step."""
+        count = tool_call_history.count(self.tool_name)
+
+        # If the tool has been used max_count_limit times, it is no longer allowed
+        if count >= self.max_count_limit:
+            return available_tools - {self.tool_name}
+
+        return available_tools
+
+
+ToolRule = Annotated[
+    Union[
+        ChildToolRule,
+        InitToolRule,
+        TerminalToolRule,
+        ConditionalToolRule,
+        ContinueToolRule,
+        MaxCountPerStepToolRule,
+        ParentToolRule,
+    ],
+    Field(discriminator="type"),
+]

mirix/schemas/usage.py

@@ -0,0 +1,31 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class MirixUsageStatistics(BaseModel):
+    """
+    Usage statistics for the agent interaction.
+
+    Attributes:
+        completion_tokens (int): The number of tokens generated by the agent.
+        prompt_tokens (int): The number of tokens in the prompt.
+        total_tokens (int): The total number of tokens processed by the agent.
+        step_count (int): The number of steps taken by the agent.
+    """
+
+    message_type: Literal["usage_statistics"] = "usage_statistics"
+    completion_tokens: int = Field(
+        0, description="The number of tokens generated by the agent."
+    )
+    prompt_tokens: int = Field(0, description="The number of tokens in the prompt.")
+    last_prompt_tokens: int = Field(
+        0, description="The number of tokens in the last prompt."
+    )
+    last_completion_tokens: int = Field(
+        0, description="The number of tokens in the last completion."
+    )
+    total_tokens: int = Field(
+        0, description="The total number of tokens processed by the agent."
+    )
+    step_count: int = Field(0, description="The number of steps taken by the agent.")

mirix/schemas/user.py

@@ -0,0 +1,67 @@
+from datetime import datetime
+from typing import Optional
+import uuid
+
+from pydantic import Field
+
+from mirix.helpers.datetime_helpers import get_utc_time
+from mirix.schemas.mirix_base import MirixBase
+from mirix.services.organization_manager import OrganizationManager
+
+
+class UserBase(MirixBase):
+    __id_prefix__ = "user"
+
+
+def _generate_user_id() -> str:
+    """Generate a random user ID."""
+    return f"user-{uuid.uuid4().hex[:8]}"
+
+
+class User(UserBase):
+    """
+    Representation of a user.
+
+    Parameters:
+        id (str): The unique identifier of the user.
+        name (str): The name of the user.
+        status (str): Whether the user is active or not.
+        created_at (datetime): The creation date of the user.
+    """
+
+    id: str = Field(
+        default_factory=_generate_user_id,
+        description="The unique identifier of the user.",
+    )
+    organization_id: Optional[str] = Field(
+        OrganizationManager.DEFAULT_ORG_ID,
+        description="The organization id of the user",
+    )
+    name: str = Field(..., description="The name of the user.")
+    status: str = Field("active", description="Whether the user is active or not.")
+    timezone: str = Field(..., description="The timezone of the user.")
+    created_at: Optional[datetime] = Field(
+        default_factory=get_utc_time, description="The creation date of the user."
+    )
+    updated_at: Optional[datetime] = Field(
+        default_factory=get_utc_time, description="The update date of the user."
+    )
+    is_deleted: bool = Field(False, description="Whether this user is deleted or not.")
+
+
+class UserCreate(UserBase):
+    id: Optional[str] = Field(None, description="The unique identifier of the user.")
+    name: str = Field(..., description="The name of the user.")
+    status: str = Field("active", description="Whether the user is active or not.")
+    timezone: str = Field(..., description="The timezone of the user.")
+    organization_id: str = Field(..., description="The organization id of the user.")
+
+
+class UserUpdate(UserBase):
+    id: str = Field(..., description="The id of the user to update.")
+    name: Optional[str] = Field(None, description="The new name of the user.")
+    status: Optional[str] = Field(None, description="The new status of the user.")
+    timezone: Optional[str] = Field(None, description="The new timezone of the user.")
+    organization_id: Optional[str] = Field(
+        None, description="The new organization id of the user."
+    )