get agent max_iter config from env-var

firstbatchxyz · Jun 6, 2024 · 4385bfb · 4385bfb
1 parent 6923e31
commit 4385bfb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/src/dria_searching_agent/config/config.py b/src/dria_searching_agent/config/config.py
@@ -7,6 +7,7 @@ def __init__(self):
 
         self.agent_model_provider = os.getenv('AGENT_MODEL_PROVIDER')
         self.agent_model = os.getenv('AGENT_MODEL')
+        self.agent_max_iter = os.getenv('AGENT_MAX_ITER', 10)
 
         self.anthropic_key = os.getenv('ANTHROPIC_KEY')
         self.openai_api_key = os.getenv('OPENAI_API_KEY')
@@ -30,6 +31,9 @@ def AGENT_MODEL_PROVIDER():
 def AGENT_MODEL():
     return config.agent_model
 
+def AGENT_MAX_ITER():
+    return config.agent_max_iter
+
 def ANTHROPIC_KEY():
     return config.anthropic_key
 

diff --git a/src/dria_searching_agent/main.py b/src/dria_searching_agent/main.py
@@ -120,6 +120,7 @@ def __create_agents(self):
         for k,v in self.agents_config.items():
             agent = Agent(
                 **v,
+                max_iter=config.AGENT_MAX_ITER(), # TODO: maybe we can get this from agents.json
                 verbose=True,
                 llm = self.__get_model(),
                 tools=[