Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import requests
 import openai
 import gradio as gr
-from openai_swarm import Agent, Swarm
+from swarm import Swarm, Agent
 
 # Fetch API keys from environment variables
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -47,40 +47,29 @@ def generate_email_content(api_key, prompt):
 
 # Function to validate the generated email for professional tone and completeness
 def validate_email(email_content):
-    # Basic validation: check if the email contains key sections (can be expanded)
     return "Why" in email_content and "How" in email_content and "What" in email_content
 
-# Define the ReAct Agent using Swarm framework
-class EmailAgent(Agent):
-    def __init__(self, user_data):
-        super().__init__()
-        self.user_data = user_data
-        self.iterations = 0
-
-    def act(self):
-        linkedin_info = fetch_linkedin_data(self.user_data['linkedin_url'])
-        company_info = fetch_company_info(self.user_data['company_url'])
-        if "error" in linkedin_info or "error" in company_info:
-            return "Error fetching data. Please check the LinkedIn and company URLs."
-        
-        prompt = structure_email(self.user_data, linkedin_info, company_info)
-        email_content = generate_email_content(OPENAI_API_KEY, prompt)
-        return email_content
-
-    def react(self, output):
-        # React based on the generated email's validity
-        if validate_email(output):
-            return output
-        else:
-            # If invalid, refine prompt or retry (up to 3 iterations)
-            self.iterations += 1
-            if self.iterations < 3:
-                refined_prompt = f"Refined: {output}"
-                return generate_email_content(OPENAI_API_KEY, refined_prompt)
-            else:
-                return "Unable to generate a valid email after 3 attempts."
-
-# Define the main function to run the agent within the Swarm
+# Define Agent B - The agent responsible for processing and refining the email
+def transfer_to_email_generation():
+    return email_agent_b
+
+# Define Agent A - The primary agent to initiate data collection and hand over to Agent B
+email_agent_a = Agent(
+    name="Data Collection Agent",
+    instructions="Collect user inputs and relevant data (LinkedIn and company details).",
+    functions=[transfer_to_email_generation]
+)
+
+# Define Agent B - The agent that structures and generates the email
+email_agent_b = Agent(
+    name="Email Generation Agent",
+    instructions="Structure the email using the 'Start with Why' model and generate professional content.",
+)
+
+# Set up the Swarm client
+client = Swarm()
+
+# Main function that interacts with the agents using the Swarm framework
 def run_agent(name, email, phone, linkedin_url, company_url, role):
     user_data = {
         "name": name,
@@ -90,15 +79,31 @@ def run_agent(name, email, phone, linkedin_url, company_url, role):
         "company_url": company_url,
         "role": role
     }
-    
-    # Initialize the swarm and add the EmailAgent
-    email_swarm = Swarm()
-    agent = EmailAgent(user_data)
-    email_swarm.add_agent(agent)
-    
-    # Run the swarm to get the final email content
-    result = email_swarm.run()
-    return result
+
+    # Step 1: Fetch LinkedIn and Company Data
+    linkedin_info = fetch_linkedin_data(linkedin_url)
+    company_info = fetch_company_info(company_url)
+    if "error" in linkedin_info or "error" in company_info:
+        return "Error fetching data. Please check the LinkedIn and company URLs."
+
+    # Step 2: Use the Swarm framework to structure and generate the email
+    messages = [{"role": "user", "content": "Initiate email generation."}]
+    response = client.run(agent=email_agent_a, messages=messages)
+
+    # Step 3: Agent B structures and refines the email using the provided data
+    prompt = structure_email(user_data, linkedin_info, company_info)
+    email_content = generate_email_content(OPENAI_API_KEY, prompt)
+
+    # Validate the generated email content using the ReAct pattern (with a max of 3 iterations)
+    for i in range(3):
+        if validate_email(email_content):
+            return email_content
+        else:
+            # Refine the prompt based on the feedback or iteration logic
+            refined_prompt = f"Refine: {prompt}"
+            email_content = generate_email_content(OPENAI_API_KEY, refined_prompt)
+
+    return "Unable to generate a valid email after 3 attempts."
 
 # Set up the Gradio interface
 final_interface = gr.Interface(