Implement exponential backoff strategy for rate limiting

praisonai/auto.py

@@ -33,6 +33,8 @@ class AutoGenerator:
         self.topic = topic
         self.agent_file = agent_file
         self.framework = framework
+        if not os.getenv("OPENAI_API_KEY"):
+            raise EnvironmentError("The OPENAI_API_KEY environment variable is not set.")
         self.client = instructor.patch(
             OpenAI(
                 base_url=self.config_list[0]['base_url'],
@@ -41,20 +43,34 @@ class AutoGenerator:
             mode=instructor.Mode.JSON,
         )
 
+    import time
+
     def generate(self):
-        response = self.client.chat.completions.create(
-            model=self.config_list[0]['model'],
-            response_model=TeamStructure,
-            max_retries=10,
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant designed to output complex team structures."},
-                {"role": "user", "content": self.get_user_content()}
-            ]
-        )
-        json_data = json.loads(response.model_dump_json())
-        self.convert_and_save(json_data)
-        full_path = os.path.abspath(self.agent_file)
-        return full_path
+        retry_delay = 1  # Start with a 1 second delay
+        max_retries = 5  # Maximum number of retries
+        retries = 0
+
+        while retries < max_retries:
+            try:
+                response = self.client.chat.completions.create(
+                    model=self.config_list[0]['model'],
+                    response_model=TeamStructure,
+                    messages=[
+                        {"role": "system", "content": "You are a helpful assistant designed to output complex team structures."},
+                        {"role": "user", "content": self.get_user_content()}
+                    ]
+                )
+                json_data = json.loads(response.model_dump_json())
+                self.convert_and_save(json_data)
+                full_path = os.path.abspath(self.agent_file)
+                return full_path
+            except openai.error.RateLimitError:
+                print(f"RateLimitError: Retrying in {retry_delay} seconds...")
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Exponential backoff
+                retries += 1
+
+        raise Exception("Max retries exceeded for API call")
 
     def convert_and_save(self, json_data):
         """Converts the provided JSON data into the desired YAML format and saves it to a file.
@@ -93,16 +109,16 @@ class AutoGenerator:
             yaml.dump(yaml_data, f, allow_unicode=True, sort_keys=False)
 
     def get_user_content(self):
-        user_content = """Generate a team structure for  \"""" + self.topic + """\" task. 
+        user_content = """Generate a team structure for  \"""" + self.topic + """\" task.
 No Input data will be provided to the team.
 The team will work in sequence. First role will pass the output to the next role, and so on.
 The last role will generate the final output.
 Think step by step.
 With maximum 3 roles, each with 1 task. Include role goals, backstories, task descriptions, and expected outputs.
 List of Available Tools: CodeDocsSearchTool, CSVSearchTool, DirectorySearchTool, DOCXSearchTool, DirectoryReadTool, FileReadTool, TXTSearchTool, JSONSearchTool, MDXSearchTool, PDFSearchTool, RagTool, ScrapeElementFromWebsiteTool, ScrapeWebsiteTool, WebsiteSearchTool, XMLSearchTool, YoutubeChannelSearchTool, YoutubeVideoSearchTool.
-Only use Available Tools. Do Not use any other tools. 
-Example Below: 
-Use below example to understand the structure of the output. 
+Only use Available Tools. Do Not use any other tools.
+Example Below:
+Use below example to understand the structure of the output.
 The final role you create should satisfy the provided task: """ + self.topic + """.
 {
 "roles": {
@@ -134,6 +150,6 @@ The final role you create should satisfy the provided task: """ + self.topic + "
         """
         return user_content
 
-    
+
 # generator = AutoGenerator(framework="crewai", topic="Create a snake game in python")
-# print(generator.generate())
+# print(generator.generate())

praisonai/cli.py

@@ -7,6 +7,7 @@ from rich import print
 from dotenv import load_dotenv
 from crewai import Agent, Task, Crew
 load_dotenv()
+print(f"Loaded environment variables: OPENAI_MODEL_NAME={os.environ.get('OPENAI_MODEL_NAME')}, OPENAI_API_KEY={os.environ.get('OPENAI_API_KEY')}, OPENAI_API_BASE={os.environ.get('OPENAI_API_BASE')}")
 import autogen
 import gradio as gr
 import argparse
@@ -41,10 +42,10 @@ class PraisonAI:
             return
         invocation_cmd = "praisonai"
         version_string = f"PraisonAI version {__version__}"
-        
+
         if args.framework:
             self.framework = args.framework
-        
+
         ui = args.ui
 
         if args.agent_file:
@@ -56,14 +57,18 @@ class PraisonAI:
         else:
             full_path = os.path.abspath(self.agent_file)
             self.agent_file = full_path
-        
+
         if args.auto or args.init:
             temp_topic = ' '.join(args.auto) if args.auto else ' '.join(args.init)
             self.topic = temp_topic
         elif self.auto or self.init:  # Use the auto attribute if args.auto is not provided
             self.topic = self.auto
-            
+
         if args.auto or self.auto:
+            if not os.environ.get('OPENAI_API_KEY'):
+                print("Error: OPENAI_API_KEY is not set. Please check your .env file.")
+                sys.exit(1)
+            print(f"Debug: OPENAI_API_KEY={os.environ.get('OPENAI_API_KEY')}")
             self.agent_file = "test.yaml"
             generator = AutoGenerator(topic=self.topic , framework=self.framework, agent_file=self.agent_file)
             self.agent_file = generator.generate()
@@ -71,19 +76,23 @@ class PraisonAI:
             result = agents_generator.generate_crew_and_kickoff()
             return result
         elif args.init or self.init:
+            if not os.environ.get('OPENAI_API_KEY'):
+                print("Error: OPENAI_API_KEY is not set. Please check your .env file.")
+                sys.exit(1)
+            print(f"Debug: OPENAI_API_KEY={os.environ.get('OPENAI_API_KEY')}")
             self.agent_file = "agents.yaml"
             generator = AutoGenerator(topic=self.topic , framework=self.framework, agent_file=self.agent_file)
             self.agent_file = generator.generate()
             print("File {} created successfully".format(self.agent_file))
             return "File {} created successfully".format(self.agent_file)
-        
+
         if ui:
             self.create_gradio_interface()
         else:
             agents_generator = AgentsGenerator(self.agent_file, self.framework, self.config_list)
             result = agents_generator.generate_crew_and_kickoff()
             return result
-            
+
     def parse_args(self):
         parser = argparse.ArgumentParser(prog="praisonai", description="praisonAI command-line interface")
         parser.add_argument("--framework", choices=["crewai", "autogen"], default="crewai", help="Specify the framework")
@@ -161,4 +170,3 @@ if __name__ == "__main__":
 #         description: Create a storyboard for the movie script about a cat in Mars.
 #         expected_output: A detailed storyboard for the movie about a cat in Mars.
 # dependencies: []
-