Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import requests
 import openai
 import gradio as gr
-from swarm import Swarm, Agent
 
 # Fetch API keys from environment variables
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -49,27 +48,42 @@ def generate_email_content(api_key, prompt):
 def validate_email(email_content):
     return "Why" in email_content and "How" in email_content and "What" in email_content
 
-# Define Agent B - The agent responsible for processing and refining the email
-def transfer_to_email_generation():
-    return email_agent_b
+# Custom Agent class to simulate behavior similar to OpenAI's Swarm framework
+class Agent:
+    def __init__(self, name, instructions, user_data):
+        self.name = name
+        self.instructions = instructions
+        self.user_data = user_data
+
+    def act(self):
+        # Agent performs its task based on instructions
+        if self.name == "Data Collection Agent":
+            linkedin_info = fetch_linkedin_data(self.user_data['linkedin_url'])
+            company_info = fetch_company_info(self.user_data['company_url'])
+            return linkedin_info, company_info
+        elif self.name == "Email Generation Agent":
+            linkedin_info, company_info = self.user_data
+            prompt = structure_email(self.user_data[0], linkedin_info, company_info)
+            email_content = generate_email_content(OPENAI_API_KEY, prompt)
+            return email_content
 
-# 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]
-)
+# Simulated Swarm class to manage agents
+class Swarm:
+    def __init__(self):
+        self.agents = []
 
-# 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.",
-)
+    def add_agent(self, agent):
+        self.agents.append(agent)
 
-# Set up the Swarm client
-client = Swarm()
+    def run(self):
+        for agent in self.agents:
+            if agent.name == "Data Collection Agent":
+                linkedin_info, company_info = agent.act()
+                if "error" in linkedin_info or "error" in company_info:
+                    return "Error fetching data. Please check the LinkedIn and company URLs."
+                return linkedin_info, company_info
 
-# Main function that interacts with the agents using the Swarm framework
+# Function that integrates the agents and manages iterations
 def run_agent(name, email, phone, linkedin_url, company_url, role):
     user_data = {
         "name": name,
@@ -80,27 +94,27 @@ def run_agent(name, email, phone, linkedin_url, company_url, role):
         "role": role
     }
 
-    # 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."
+    # Create a Swarm and add the Data Collection Agent
+    email_swarm = Swarm()
+    data_collection_agent = Agent("Data Collection Agent", "Collect user inputs and relevant data", user_data)
+    email_swarm.add_agent(data_collection_agent)
 
-    # 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)
+    # Get data from the Data Collection Agent
+    linkedin_info, company_info = email_swarm.run()
+    if isinstance(linkedin_info, str):  # If an error message is returned
+        return linkedin_info
 
-    # 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)
+    # Pass the collected data to the Email Generation Agent
+    email_agent = Agent("Email Generation Agent", "Generate the email content", (user_data, linkedin_info, company_info))
+    email_content = email_agent.act()
 
-    # Validate the generated email content using the ReAct pattern (with a max of 3 iterations)
+    # Validate and refine the email using a ReAct pattern with a maximum 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}"
+            refined_prompt = f"Refine: {structure_email(user_data, linkedin_info, company_info)}"
             email_content = generate_email_content(OPENAI_API_KEY, refined_prompt)
 
     return "Unable to generate a valid email after 3 attempts."