Add support for max_tool_count and max iterations from yaml

[vizsatiz]

• Currently the agent will stop looping when it reaches max_tool_count, and break the flow. Right now the max loop count is exposed in constructor of Agent Class, but its not configurable from YAML.
• Similarly the execution will break when arium reaches max_iterations, this is hardcoded to 20 as of now. This should be exposed over arium interface & also in YAML

[harikapadia999]

Implementation Proposal: YAML-Configurable Agent Limits

Great enhancement request! Making max_tool_count and max_iterations configurable via YAML will significantly improve agent control and prevent runaway loops.

Proposed YAML Schema:

agents:
  - name: "customer_support_agent"
    model: "gpt-4"
    limits:
      max_tool_count: 10        # Maximum number of tool calls per session
      max_iterations: 5         # Maximum conversation turns
      timeout_seconds: 300      # Optional: Overall timeout
      max_tokens_per_call: 4096 # Optional: Token limit per call
    tools:
      - search_knowledge_base
      - create_ticket

Implementation Approach:

1. Update Agent Configuration Schema

# flo_ai/config/agent_schema.py
from pydantic import BaseModel, Field

class AgentLimits(BaseModel):
    max_tool_count: int = Field(default=10, ge=1, le=100)
    max_iterations: int = Field(default=5, ge=1, le=50)
    timeout_seconds: Optional[int] = Field(default=None, ge=1)
    max_tokens_per_call: Optional[int] = Field(default=None, ge=1)

class AgentConfig(BaseModel):
    name: str
    model: str
    limits: AgentLimits = Field(default_factory=AgentLimits)
    tools: List[str]

2. Update Agent Class Constructor

# flo_ai/agent/agent.py
class Agent:
    def __init__(
        self,
        name: str,
        model: str,
        tools: List[Tool],
        limits: Optional[AgentLimits] = None,
        **kwargs
    ):
        self.name = name
        self.model = model
        self.tools = tools
        
        # Use provided limits or defaults
        self.limits = limits or AgentLimits()
        self.max_tool_count = self.limits.max_tool_count
        self.max_iterations = self.limits.max_iterations
        self.timeout_seconds = self.limits.timeout_seconds

3. Enforce Limits in Agent Loop

def run(self, user_input: str) -> str:
    start_time = time.time()
    iteration = 0
    tool_call_count = 0
    
    while iteration < self.max_iterations:
        # Check timeout
        if self.timeout_seconds and (time.time() - start_time) > self.timeout_seconds:
            raise TimeoutError(f"Agent exceeded timeout of {self.timeout_seconds}s")
        
        response = self.llm.generate(...)
        
        # Check tool call limit
        if response.tool_calls:
            tool_call_count += len(response.tool_calls)
            if tool_call_count > self.max_tool_count:
                logger.warning(f"Max tool count ({self.max_tool_count}) reached")
                return self._generate_limit_exceeded_response()
        
        iteration += 1
    
    logger.warning(f"Max iterations ({self.max_iterations}) reached")
    return self._generate_limit_exceeded_response()

4. YAML Loader Integration

# flo_ai/config/loader.py
def load_agent_from_yaml(yaml_path: str) -> Agent:
    with open(yaml_path, 'r') as f:
        config = yaml.safe_load(f)
    
    agent_config = AgentConfig(**config['agents'][0])
    
    return Agent(
        name=agent_config.name,
        model=agent_config.model,
        tools=load_tools(agent_config.tools),
        limits=agent_config.limits
    )

Backward Compatibility:

Ensure existing agents without YAML limits still work:

# Default limits if not specified
DEFAULT_LIMITS = AgentLimits(
    max_tool_count=10,
    max_iterations=5
)

Testing Strategy:

def test_yaml_limits():
    yaml_content = """
    agents:
      - name: test_agent
        model: gpt-4
        limits:
          max_tool_count: 3
          max_iterations: 2
        tools: []
    """
    
    agent = load_agent_from_yaml(yaml_content)
    assert agent.max_tool_count == 3
    assert agent.max_iterations == 2

def test_limit_enforcement():
    agent = Agent(
        name="test",
        model="gpt-4",
        tools=[],
        limits=AgentLimits(max_tool_count=2)
    )
    
    # Simulate exceeding limit
    with pytest.raises(LimitExceededError):
        agent.run_with_excessive_tool_calls()

Documentation Update:

Add to TOOLS.md:

## Agent Limits Configuration

Control agent behavior with configurable limits:

\`\`\`yaml
agents:
  - name: "my_agent"
    limits:
      max_tool_count: 10      # Stop after 10 tool calls
      max_iterations: 5       # Stop after 5 conversation turns
      timeout_seconds: 300    # Stop after 5 minutes
\`\`\`

**Best Practices:**
- Set `max_tool_count` based on expected workflow complexity
- Use `max_iterations` to prevent infinite loops
- Add `timeout_seconds` for long-running agents
\`\`\`

### Benefits:

1. ✅ Prevents runaway agents (addresses issue #172)
2. ✅ Configurable per-agent without code changes
3. ✅ Backward compatible with existing agents
4. ✅ Clear error messages when limits are exceeded

Would love to implement this! Let me know if you'd like me to create a PR with this approach.

[Keesan12]

YAML-configurable limits are a good control point, but I would avoid making every limit breach throw the same kind of broken-flow error.

For agent loops, the result should distinguish:

• hard failure: the run cannot produce a useful answer
• graceful cap: the agent hit max_iterations or max_tool_count and returned best effort
• policy denial: the next tool call was not admitted because the configured cap was already reached

That distinction matters for callers. A workflow might retry after a transient tool error, but it should not blindly retry after a policy denial with the same prompt and same limits. Adding a structured stop reason to the agent result would make the YAML settings much easier to operate.