130 lines
5.9 KiB
Python
130 lines
5.9 KiB
Python
from typing import List, Optional, Set, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from letta.schemas.enums import ToolRuleType
|
|
from letta.schemas.tool_rule import (
|
|
BaseToolRule,
|
|
ChildToolRule,
|
|
ConditionalToolRule,
|
|
ContinueToolRule,
|
|
InitToolRule,
|
|
MaxCountPerStepToolRule,
|
|
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."
|
|
)
|
|
continue_tool_rules: List[ContinueToolRule] = Field(
|
|
default_factory=list, description="Continue tool rules to be used to continue tool execution."
|
|
)
|
|
# TODO: This should be renamed?
|
|
# TODO: These are tools that control the set of allowed functions in the next turn
|
|
child_based_tool_rules: List[Union[ChildToolRule, ConditionalToolRule, MaxCountPerStepToolRule]] = 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."
|
|
)
|
|
tool_call_history: List[str] = Field(default_factory=list, description="History of tool calls, 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.child_based_tool_rules.append(rule)
|
|
elif rule.type == ToolRuleType.conditional:
|
|
assert isinstance(rule, ConditionalToolRule)
|
|
self.validate_conditional_tool(rule)
|
|
self.child_based_tool_rules.append(rule)
|
|
elif rule.type == ToolRuleType.exit_loop:
|
|
assert isinstance(rule, TerminalToolRule)
|
|
self.terminal_tool_rules.append(rule)
|
|
elif rule.type == ToolRuleType.continue_loop:
|
|
assert isinstance(rule, ContinueToolRule)
|
|
self.continue_tool_rules.append(rule)
|
|
elif rule.type == ToolRuleType.max_count_per_step:
|
|
assert isinstance(rule, MaxCountPerStepToolRule)
|
|
self.child_based_tool_rules.append(rule)
|
|
|
|
def update_tool_usage(self, tool_name: str):
|
|
"""Update the internal state to track tool call history."""
|
|
self.tool_call_history.append(tool_name)
|
|
|
|
def clear_tool_history(self):
|
|
"""Clear the history of tool calls."""
|
|
self.tool_call_history.clear()
|
|
|
|
def get_allowed_tool_names(
|
|
self, available_tools: Set[str], 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."""
|
|
# TODO: This piece of code here is quite ugly and deserves a refactor
|
|
# TODO: There's some weird logic encoded here:
|
|
# TODO: -> This only takes into consideration Init, and a set of Child/Conditional/MaxSteps tool rules
|
|
# TODO: -> Init tool rules outputs are treated additively, Child/Conditional/MaxSteps are intersection based
|
|
# TODO: -> Tool rules should probably be refactored to take in a set of tool names?
|
|
# If no tool has been called yet, return InitToolRules additively
|
|
if not self.tool_call_history:
|
|
if self.init_tool_rules:
|
|
# If there are init tool rules, only return those defined in the init tool rules
|
|
return [rule.tool_name for rule in self.init_tool_rules]
|
|
else:
|
|
# Otherwise, return all the available tools
|
|
return list(available_tools)
|
|
else:
|
|
# Collect valid tools from all child-based rules
|
|
valid_tool_sets = [
|
|
rule.get_valid_tools(self.tool_call_history, available_tools, last_function_response)
|
|
for rule in self.child_based_tool_rules
|
|
]
|
|
|
|
# Compute intersection of all valid tool sets
|
|
final_allowed_tools = set.intersection(*valid_tool_sets) if valid_tool_sets else available_tools
|
|
|
|
if error_on_empty and not final_allowed_tools:
|
|
raise ValueError("No valid tools found based on tool rules.")
|
|
|
|
return list(final_allowed_tools)
|
|
|
|
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.child_based_tool_rules)
|
|
|
|
def is_continue_tool(self, tool_name):
|
|
"""Check if the tool is defined as a continue tool in the tool rules."""
|
|
return any(rule.tool_name == tool_name for rule in self.continue_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
|