letta-nightly 0.5.4.dev20241126104249__py3-none-any.whl → 0.5.4.dev20241128000451__py3-none-any.whl

letta/schemas/agent.py

@@ -2,15 +2,18 @@ from datetime import datetime
 from enum import Enum
 from typing import Dict, List, Optional
 
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, Field, field_validator
 
+from letta.schemas.block import CreateBlock
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.letta_base import LettaBase
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.memory import Memory
 from letta.schemas.message import Message
 from letta.schemas.openai.chat_completion_response import UsageStatistics
-from letta.schemas.tool_rule import BaseToolRule
+from letta.schemas.source import Source
+from letta.schemas.tool import Tool
+from letta.schemas.tool_rule import ToolRule
 
 
 class BaseAgent(LettaBase, validate_assignment=True):
@@ -32,23 +35,8 @@ class AgentType(str, Enum):
     o1_agent = "o1_agent"
 
 
-class AgentState(BaseAgent, 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.
-        message_ids (List[str]): The ids of the messages in the agent's in-context memory.
-        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.
-
-    """
-
+class PersistedAgentState(BaseAgent, validate_assignment=True):
+    # NOTE: this has been changed to represent the data stored in the ORM, NOT what is passed around internally or returned to the user
     id: str = BaseAgent.generate_id_field()
     name: str = Field(..., description="The name of the agent.")
     created_at: datetime = Field(..., description="The datetime the agent was created.", default_factory=datetime.now)
@@ -56,16 +44,12 @@ class AgentState(BaseAgent, validate_assignment=True):
     # in-context memory
     message_ids: Optional[List[str]] = Field(default=None, description="The ids of the messages in the agent's in-context memory.")
 
-    memory: Memory = Field(default_factory=Memory, description="The in-context memory of the agent.")
-
     # tools
-    tools: List[str] = Field(..., description="The tools used by the agent.")
+    # TODO: move to ORM mapping
+    tool_names: List[str] = Field(..., description="The tools used by the agent.")
 
     # tool rules
-    tool_rules: Optional[List[BaseToolRule]] = Field(default=None, description="The list of tool rules.")
-
-    # tags
-    tags: Optional[List[str]] = Field(None, description="The tags associated with the agent.")
+    tool_rules: Optional[List[ToolRule]] = Field(default=None, description="The list of tool rules.")
 
     # system prompt
     system: str = Field(..., description="The system prompt used by the agent.")
@@ -77,40 +61,62 @@ class AgentState(BaseAgent, validate_assignment=True):
     llm_config: LLMConfig = Field(..., description="The LLM configuration used by the agent.")
     embedding_config: EmbeddingConfig = Field(..., description="The embedding configuration used by the agent.")
 
-    def __init__(self, **data):
-        super().__init__(**data)
-        self._internal_memory = self.memory
+    class Config:
+        arbitrary_types_allowed = True
+        validate_assignment = True
 
-    @model_validator(mode="after")
-    def verify_memory_type(self):
-        try:
-            assert isinstance(self.memory, Memory)
-        except Exception as e:
-            raise e
-        return self
 
-    @property
-    def memory(self) -> Memory:
-        return self._internal_memory
+class AgentState(PersistedAgentState):
+    """
+    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.
+        message_ids (List[str]): The ids of the messages in the agent's in-context memory.
+        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.
 
-    @memory.setter
-    def memory(self, value):
-        if not isinstance(value, Memory):
-            raise TypeError(f"Expected Memory, got {type(value).__name__}")
-        self._internal_memory = value
+    """
 
-    class Config:
-        arbitrary_types_allowed = True
-        validate_assignment = True
+    # NOTE: this is what is returned to the client and also what is used to initialize `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
+    memory: Memory = Field(..., description="The in-context memory of the agent.")
+    tools: List[Tool] = Field(..., description="The tools used by the agent.")
+    sources: List[Source] = Field(..., description="The sources used by the agent.")
+    tags: List[str] = Field(..., description="The tags associated with the agent.")
+    # TODO: add in context message objects
 
+    def to_persisted_agent_state(self) -> PersistedAgentState:
+        # turn back into persisted agent
+        data = self.model_dump()
+        del data["memory"]
+        del data["tools"]
+        del data["sources"]
+        del data["tags"]
+        return PersistedAgentState(**data)
 
-class CreateAgent(BaseAgent):
+
+class CreateAgent(BaseAgent):  #
     # all optional as server can generate defaults
     name: Optional[str] = Field(None, description="The name of the agent.")
     message_ids: Optional[List[str]] = Field(None, description="The ids of the messages in the agent's in-context memory.")
-    memory: Optional[Memory] = Field(None, description="The in-context memory of the agent.")
+
+    # memory creation
+    memory_blocks: List[CreateBlock] = Field(
+        # [CreateHuman(), CreatePersona()], description="The blocks to create in the agent's in-context memory."
+        ...,
+        description="The blocks to create in the agent's in-context memory.",
+    )
+
     tools: Optional[List[str]] = Field(None, description="The tools used by the agent.")
-    tool_rules: Optional[List[BaseToolRule]] = Field(None, description="The tool rules governing the agent.")
+    tool_rules: Optional[List[ToolRule]] = Field(None, description="The tool rules governing the agent.")
     tags: Optional[List[str]] = Field(None, description="The tags associated with the agent.")
     system: Optional[str] = Field(None, description="The system prompt used by the agent.")
     agent_type: Optional[AgentType] = Field(None, description="The type of agent.")
@@ -151,7 +157,7 @@ class CreateAgent(BaseAgent):
 class UpdateAgentState(BaseAgent):
     id: str = Field(..., description="The id of the agent.")
     name: Optional[str] = Field(None, description="The name of the agent.")
-    tools: Optional[List[str]] = Field(None, description="The tools used by the agent.")
+    tool_names: Optional[List[str]] = Field(None, description="The tools used by the agent.")
     tags: Optional[List[str]] = Field(None, description="The tags associated with the agent.")
     system: Optional[str] = Field(None, description="The system prompt used by the agent.")
     llm_config: Optional[LLMConfig] = Field(None, description="The LLM configuration used by the agent.")
@@ -159,7 +165,6 @@ class UpdateAgentState(BaseAgent):
 
     # TODO: determine if these should be editable via this schema?
     message_ids: Optional[List[str]] = Field(None, description="The ids of the messages in the agent's in-context memory.")
-    memory: Optional[Memory] = Field(None, description="The in-context memory of the agent.")
 
 
 class AgentStepResponse(BaseModel):

letta/schemas/block.py

@@ -28,17 +28,20 @@ class BaseBlock(LettaBase, validate_assignment=True):
     description: Optional[str] = Field(None, description="Description of the block.")
     metadata_: Optional[dict] = Field({}, description="Metadata of the block.")
 
+    # def __len__(self):
+    #     return len(self.value)
+
+    class Config:
+        extra = "ignore"  # Ignores extra fields
+
     @model_validator(mode="after")
     def verify_char_limit(self) -> Self:
-        if len(self.value) > self.limit:
+        if self.value and len(self.value) > self.limit:
             error_msg = f"Edit failed: Exceeds {self.limit} character limit (requested {len(self.value)}) - {str(self)}."
             raise ValueError(error_msg)
 
         return self
 
-    # def __len__(self):
-    #     return len(self.value)
-
     def __setattr__(self, name, value):
         """Run validation if self.value is updated"""
         super().__setattr__(name, value)
@@ -46,9 +49,6 @@ class BaseBlock(LettaBase, validate_assignment=True):
             # run validation
             self.__class__.model_validate(self.model_dump(exclude_unset=True))
 
-    class Config:
-        extra = "ignore"  # Ignores extra fields
-
 
 class Block(BaseBlock):
     """
@@ -88,11 +88,11 @@ class Persona(Block):
     label: str = "persona"
 
 
-class BlockCreate(BaseBlock):
-    """Create a block"""
-
-    is_template: bool = True
-    label: str = Field(..., description="Label of the block.")
+# class CreateBlock(BaseBlock):
+#    """Create a block"""
+#
+#    is_template: bool = True
+#    label: str = Field(..., description="Label of the block.")
 
 
 class BlockLabelUpdate(BaseModel):
@@ -102,16 +102,16 @@ class BlockLabelUpdate(BaseModel):
     new_label: str = Field(..., description="New label of the block.")
 
 
-class CreatePersona(BlockCreate):
-    """Create a persona block"""
-
-    label: str = "persona"
-
-
-class CreateHuman(BlockCreate):
-    """Create a human block"""
-
-    label: str = "human"
+# class CreatePersona(CreateBlock):
+#    """Create a persona block"""
+#
+#    label: str = "persona"
+#
+#
+# class CreateHuman(CreateBlock):
+#    """Create a human block"""
+#
+#    label: str = "human"
 
 
 class BlockUpdate(BaseBlock):
@@ -131,13 +131,57 @@ class BlockLimitUpdate(BaseModel):
     limit: int = Field(..., description="New limit of the block.")
 
 
-class UpdatePersona(BlockUpdate):
-    """Update a persona block"""
+# class UpdatePersona(BlockUpdate):
+#    """Update a persona block"""
+#
+#    label: str = "persona"
+#
+#
+# class UpdateHuman(BlockUpdate):
+#    """Update a human block"""
+#
+#    label: str = "human"
+
+
+class CreateBlock(BaseBlock):
+    """Create a block"""
+
+    label: str = Field(..., description="Label of the block.")
+    limit: int = Field(2000, description="Character limit of the block.")
+    value: str = Field(..., description="Value of the block.")
+
+    # block templates
+    is_template: bool = False
+    template_name: Optional[str] = Field(None, description="Name of the block if it is a template.", alias="name")
+
+
+class CreateHuman(CreateBlock):
+    """Create a human block"""
+
+    label: str = "human"
+
+
+class CreatePersona(CreateBlock):
+    """Create a persona block"""
 
     label: str = "persona"
 
 
-class UpdateHuman(BlockUpdate):
-    """Update a human block"""
+class CreateBlockTemplate(CreateBlock):
+    """Create a block template"""
+
+    is_template: bool = True
+
+
+class CreateHumanBlockTemplate(CreateHuman):
+    """Create a human block template"""
 
+    is_template: bool = True
     label: str = "human"
+
+
+class CreatePersonaBlockTemplate(CreatePersona):
+    """Create a persona block template"""
+
+    is_template: bool = True
+    label: str = "persona"

letta/schemas/enums.py

@@ -33,3 +33,17 @@ class MessageStreamStatus(str, Enum):
     done_generation = "[DONE_GEN]"
     done_step = "[DONE_STEP]"
     done = "[DONE]"
+
+
+class ToolRuleType(str, Enum):
+    """
+    Type of tool rule.
+    """
+
+    # note: some of these should be renamed when we do the data migration
+
+    run_first = "InitToolRule"
+    exit_loop = "TerminalToolRule"  # reasoning loop should exit
+    continue_loop = "continue_loop"  # reasoning loop should continue
+    constrain_child_tools = "ToolRule"
+    require_parent_tools = "require_parent_tools"

letta/schemas/letta_base.py

@@ -63,7 +63,7 @@ class LettaBase(BaseModel):
     @classmethod
     def _id_example(cls, prefix: str):
         """generates an example id for a given prefix"""
-        return [prefix + "-123e4567-e89b-12d3-a456-426614174000"]
+        return f"{prefix}-123e4567-e89b-12d3-a456-426614174000"
 
     @classmethod
     def _id_description(cls, prefix: str):

letta/schemas/letta_request.py

@@ -8,33 +8,21 @@ from letta.schemas.message import Message, MessageCreate
 
 class LettaRequest(BaseModel):
     messages: Union[List[MessageCreate], List[Message]] = Field(..., description="The messages to be sent to the agent.")
-    run_async: bool = Field(default=False, description="Whether to asynchronously send the messages to the agent.")  # TODO: implement
-
-    stream_steps: bool = Field(
-        default=False, description="Flag to determine if the response should be streamed. Set to True for streaming agent steps."
-    )
-    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).",
-    )
-
-    return_message_object: bool = Field(
-        default=False,
-        description="Set True to return the raw Message object. Set False to return the Message in the format of the Letta API.",
-    )
 
     # Flags to support the use of AssistantMessage message types
 
-    use_assistant_message: bool = Field(
-        default=False,
-        description="[Only applicable if return_message_object is False] If true, returns AssistantMessage objects when the agent calls a designated message tool. If false, return FunctionCallMessage objects for all tool calls.",
-    )
-
-    assistant_message_function_name: str = Field(
+    assistant_message_tool_name: str = Field(
         default=DEFAULT_MESSAGE_TOOL,
-        description="[Only applicable if use_assistant_message is True] The name of the designated message tool.",
+        description="The name of the designated message tool.",
     )
-    assistant_message_function_kwarg: str = Field(
+    assistant_message_tool_kwarg: str = Field(
         default=DEFAULT_MESSAGE_TOOL_KWARG,
-        description="[Only applicable if use_assistant_message is True] The name of the message argument in the designated message tool.",
+        description="The name of the message argument in the designated message tool.",
+    )
+
+
+class LettaStreamingRequest(LettaRequest):
+    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).",
     )

letta/schemas/letta_response.py

@@ -7,7 +7,6 @@ from pydantic import BaseModel, Field
 
 from letta.schemas.enums import MessageStreamStatus
 from letta.schemas.letta_message import LettaMessage, LettaMessageUnion
-from letta.schemas.message import Message
 from letta.schemas.usage import LettaUsageStatistics
 from letta.utils import json_dumps
 
@@ -24,7 +23,7 @@ class LettaResponse(BaseModel):
         usage (LettaUsageStatistics): The usage statistics
     """
 
-    messages: Union[List[Message], List[LettaMessageUnion]] = Field(..., description="The messages returned by the agent.")
+    messages: List[LettaMessageUnion] = Field(..., description="The messages returned by the agent.")
     usage: LettaUsageStatistics = Field(..., description="The usage statistics of the agent.")
 
     def __str__(self):

letta/schemas/memory.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, List, Optional
 
 from jinja2 import Template, TemplateSyntaxError
 from pydantic import BaseModel, Field
@@ -55,19 +55,16 @@ class ContextWindowOverview(BaseModel):
 class Memory(BaseModel, validate_assignment=True):
     """
 
-    Represents the in-context memory of the agent. This includes both the `Block` objects (labelled by sections), as well as tools to edit the blocks.
-
-    Attributes:
-        memory (Dict[str, Block]): Mapping from memory block section to memory block.
+    Represents the in-context memory (i.e. Core memory) of the agent. This includes both the `Block` objects (labelled by sections), as well as tools to edit the blocks.
 
     """
 
-    # Memory.memory is a dict mapping from memory block label to memory block.
-    memory: Dict[str, Block] = Field(default_factory=dict, description="Mapping from memory block section to memory block.")
+    # Memory.block contains the list of memory blocks in the core memory
+    blocks: List[Block] = Field(..., description="Memory blocks contained in the agent's in-context memory")
 
     # Memory.template is a Jinja2 template for compiling memory module into a prompt string.
     prompt_template: str = Field(
-        default="{% for block in memory.values() %}"
+        default="{% for block in blocks %}"
         '<{{ block.label }} characters="{{ block.value|length }}/{{ block.limit }}">\n'
         "{{ block.value }}\n"
         "</{{ block.label }}>"
@@ -90,7 +87,7 @@ class Memory(BaseModel, validate_assignment=True):
             Template(prompt_template)
 
             # Validate compatibility with current memory structure
-            test_render = Template(prompt_template).render(memory=self.memory)
+            test_render = Template(prompt_template).render(blocks=self.blocks)
 
             # If we get here, the template is valid and compatible
             self.prompt_template = prompt_template
@@ -99,107 +96,49 @@ class Memory(BaseModel, validate_assignment=True):
         except Exception as e:
             raise ValueError(f"Prompt template is not compatible with current memory structure: {str(e)}")
 
-    @classmethod
-    def load(cls, state: dict):
-        """Load memory from dictionary object"""
-        obj = cls()
-        if len(state.keys()) == 2 and "memory" in state and "prompt_template" in state:
-            # New format
-            obj.prompt_template = state["prompt_template"]
-            for key, value in state["memory"].items():
-                # TODO: This is migration code, please take a look at a later time to get rid of this
-                if "name" in value:
-                    value["template_name"] = value["name"]
-                    value.pop("name")
-                obj.memory[key] = Block(**value)
-        else:
-            # Old format (pre-template)
-            for key, value in state.items():
-                obj.memory[key] = Block(**value)
-        return obj
-
     def compile(self) -> str:
         """Generate a string representation of the memory in-context using the Jinja2 template"""
         template = Template(self.prompt_template)
-        return template.render(memory=self.memory)
-
-    def to_dict(self):
-        """Convert to dictionary representation"""
-        return {
-            "memory": {key: value.model_dump() for key, value in self.memory.items()},
-            "prompt_template": self.prompt_template,
-        }
-
-    def to_flat_dict(self):
-        """Convert to a dictionary that maps directly from block label to values"""
-        return {k: v.value for k, v in self.memory.items() if v is not None}
+        return template.render(blocks=self.blocks)
 
     def list_block_labels(self) -> List[str]:
         """Return a list of the block names held inside the memory object"""
-        return list(self.memory.keys())
+        # return list(self.memory.keys())
+        return [block.label for block in self.blocks]
 
     # TODO: these should actually be label, not name
     def get_block(self, label: str) -> Block:
         """Correct way to index into the memory.memory field, returns a Block"""
-        if label not in self.memory:
-            raise KeyError(f"Block field {label} does not exist (available sections = {', '.join(list(self.memory.keys()))})")
-        else:
-            return self.memory[label]
+        keys = []
+        for block in self.blocks:
+            if block.label == label:
+                return block
+            keys.append(block.label)
+        raise KeyError(f"Block field {label} does not exist (available sections = {', '.join(keys)})")
 
     def get_blocks(self) -> List[Block]:
         """Return a list of the blocks held inside the memory object"""
-        return list(self.memory.values())
-
-    def link_block(self, block: Block, override: Optional[bool] = False):
-        """Link a new block to the memory object"""
-        if not isinstance(block, Block):
-            raise ValueError(f"Param block must be type Block (not {type(block)})")
-        if not override and block.label in self.memory:
-            raise ValueError(f"Block with label {block.label} already exists")
-
-        self.memory[block.label] = block
+        # return list(self.memory.values())
+        return self.blocks
 
-    def unlink_block(self, block_label: str) -> Block:
-        """Unlink a block from the memory object"""
-        if block_label not in self.memory:
-            raise ValueError(f"Block with label {block_label} does not exist")
-
-        return self.memory.pop(block_label)
+    def set_block(self, block: Block):
+        """Set a block in the memory object"""
+        for i, b in enumerate(self.blocks):
+            if b.label == block.label:
+                self.blocks[i] = block
+                return
+        self.blocks.append(block)
 
     def update_block_value(self, label: str, value: str):
         """Update the value of a block"""
-        if label not in self.memory:
-            raise ValueError(f"Block with label {label} does not exist")
         if not isinstance(value, str):
             raise ValueError(f"Provided value must be a string")
 
-        self.memory[label].value = value
-
-    def update_block_label(self, current_label: str, new_label: str):
-        """Update the label of a block"""
-        if current_label not in self.memory:
-            raise ValueError(f"Block with label {current_label} does not exist")
-        if not isinstance(new_label, str):
-            raise ValueError(f"Provided new label must be a string")
-
-        # First change the label of the block
-        self.memory[current_label].label = new_label
-
-        # Then swap the block to the new label
-        self.memory[new_label] = self.memory.pop(current_label)
-
-    def update_block_limit(self, label: str, limit: int):
-        """Update the limit of a block"""
-        if label not in self.memory:
-            raise ValueError(f"Block with label {label} does not exist")
-        if not isinstance(limit, int):
-            raise ValueError(f"Provided limit must be an integer")
-
-        # Check to make sure the new limit is greater than the current length of the block
-        if len(self.memory[label].value) > limit:
-            raise ValueError(f"New limit {limit} is less than the current length of the block {len(self.memory[label].value)}")
-
-        self.memory[label].limit = limit
+        for block in self.blocks:
+            if block.label == label:
+                block.value = value
+                return
+        raise ValueError(f"Block with label {label} does not exist")
 
 
 # TODO: ideally this is refactored into ChatMemory and the subclasses are given more specific names.
@@ -222,13 +161,7 @@ class BasicBlockMemory(Memory):
         Args:
             blocks (List[Block]): List of blocks to be linked to the memory object.
         """
-        super().__init__()
-        for block in blocks:
-            # TODO: centralize these internal schema validations
-            # assert block.name is not None and block.name != "", "each existing chat block must have a name"
-            # self.link_block(name=block.name, block=block)
-            assert block.label is not None and block.label != "", "each existing chat block must have a name"
-            self.link_block(block=block)
+        super().__init__(blocks=blocks)
 
     def core_memory_append(agent_state: "AgentState", label: str, content: str) -> Optional[str]:  # type: ignore
         """
@@ -280,9 +213,7 @@ class ChatMemory(BasicBlockMemory):
             human (str): The starter value for the human block.
             limit (int): The character limit for each block.
         """
-        super().__init__()
-        self.link_block(block=Block(value=persona, limit=limit, label="persona"))
-        self.link_block(block=Block(value=human, limit=limit, label="human"))
+        super().__init__(blocks=[Block(value=persona, limit=limit, label="persona"), Block(value=human, limit=limit, label="human")])
 
 
 class UpdateMemory(BaseModel):

letta/schemas/message.py

@@ -134,8 +134,8 @@ class Message(BaseMessage):
     def to_letta_message(
         self,
         assistant_message: bool = False,
-        assistant_message_function_name: str = DEFAULT_MESSAGE_TOOL,
-        assistant_message_function_kwarg: str = DEFAULT_MESSAGE_TOOL_KWARG,
+        assistant_message_tool_name: str = DEFAULT_MESSAGE_TOOL,
+        assistant_message_tool_kwarg: str = DEFAULT_MESSAGE_TOOL_KWARG,
     ) -> List[LettaMessage]:
         """Convert message object (in DB format) to the style used by the original Letta API"""
 
@@ -156,7 +156,7 @@ class Message(BaseMessage):
                 for tool_call in self.tool_calls:
                     # If we're supporting using assistant message,
                     # then we want to treat certain function calls as a special case
-                    if assistant_message and tool_call.function.name == assistant_message_function_name:
+                    if 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 = json.loads(tool_call.function.arguments)

letta/schemas/tool_rule.py

@@ -1,21 +1,24 @@
-from typing import List
+from typing import List, Union
 
 from pydantic import Field
 
+from letta.schemas.enums import ToolRuleType
 from letta.schemas.letta_base import LettaBase
 
 
 class BaseToolRule(LettaBase):
     __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
 
 
-class ToolRule(BaseToolRule):
+class ChildToolRule(BaseToolRule):
     """
     A ToolRule represents a tool that can be invoked by the agent.
     """
 
-    type: str = Field("ToolRule")
+    # type: str = Field("ToolRule")
+    type: ToolRuleType = ToolRuleType.constrain_child_tools
     children: List[str] = Field(..., description="The children tools that can be invoked.")
 
 
@@ -24,7 +27,8 @@ class InitToolRule(BaseToolRule):
     Represents the initial tool rule configuration.
     """
 
-    type: str = Field("InitToolRule")
+    # type: str = Field("InitToolRule")
+    type: ToolRuleType = ToolRuleType.run_first
 
 
 class TerminalToolRule(BaseToolRule):
@@ -32,4 +36,8 @@ class TerminalToolRule(BaseToolRule):
     Represents a terminal tool rule configuration where if this tool gets called, it must end the agent loop.
     """
 
-    type: str = Field("TerminalToolRule")
+    # type: str = Field("TerminalToolRule")
+    type: ToolRuleType = ToolRuleType.exit_loop
+
+
+ToolRule = Union[ChildToolRule, InitToolRule, TerminalToolRule]