Add configurable execution limits for beeai

Author: lbarcziova

beeai/agents/base_agent.py

@@ -1,8 +1,10 @@
+import os
 from abc import ABC, abstractmethod
 from typing import TypeVar
 
 from pydantic import BaseModel
 
+from beeai_framework.agents import AgentExecutionConfig
 from beeai_framework.agents.experimental import RequirementAgent, RequirementAgentRunOutput
 from beeai_framework.template import PromptTemplate, PromptTemplateInput
 
@@ -31,9 +33,18 @@ def _render_prompt(self, input: TInputSchema) -> str:
         return template.render(input)
 
     async def _run_with_schema(self, input: TInputSchema) -> TOutputSchema:
+        max_retries_per_step = int(os.getenv("BEEAI_MAX_RETRIES_PER_STEP", 5))
+        total_max_retries = int(os.getenv("BEEAI_TOTAL_MAX_RETRIES", 10))
+        max_iterations = int(os.getenv("BEEAI_MAX_ITERATIONS", 100))
+
         response = await self.run(
             prompt=self._render_prompt(input),
             expected_output=self.output_schema,
+            execution=AgentExecutionConfig(
+                max_retries_per_step=max_retries_per_step,
+                total_max_retries=total_max_retries,
+                max_iterations=max_iterations,
+            ),
         )
         return self.output_schema.model_validate_json(response.result.text)
 

beeai/templates/beeai-agent.env

@@ -5,5 +5,10 @@ CHAT_MODEL=
 #GEMINI_API_KEY=
 #ANTHROPIC_API_KEY=
 
+# BeeAI Framework Configuration
+BEEAI_MAX_ITERATIONS=100
+BEEAI_MAX_RETRIES_PER_STEP=5
+BEEAI_TOTAL_MAX_RETRIES=10
+
 GITLAB_USER=
 GITLAB_TOKEN=