Update app.py

app.py

@@ -11,36 +11,43 @@ generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 email_history = []
 
 # Function to generate email, rate it, and suggest changes
-def generate_email(subject, previous_draft=None, step_count=1):
-    # Generate initial email or improved version based on previous draft
-    if step_count == 1:
-        prompt = f"Write a formal and professional email on the topic: '{subject}'. Keep it concise, relevant, and polite."
-    else:
-        prompt = (f"Here is the current draft of an email about '{subject}':\n\n"
-                  f"{previous_draft}\n\n"
-                  f"Please improve this draft based on the suggestions given: \n\n"
-                  f"Suggestions: {suggestion}\n\n"
-                  f"This is step {step_count}.")
-    
-    response = generator(prompt, max_length=300, num_return_sequences=1)
-    generated_email = response[0]['generated_text']
+def generate_email(subject):
+    step_count = 1
+    previous_draft = None
+    score = 0
+
+    while score < 7:
+        # Generate initial email or improved version based on previous draft
+        if step_count == 1:
+            prompt = f"Write a formal and professional email on the topic: '{subject}'. Keep it concise, relevant, and polite."
+        else:
+            prompt = (f"Here is the current draft of an email about '{subject}':\n\n"
+                      f"{previous_draft}\n\n"
+                      f"This is step {step_count}. Improve this email.")
+
+        response = generator(prompt, max_length=300, num_return_sequences=1)
+        generated_email = response[0]['generated_text']
 
-    # Rate the email and suggest improvements
-    score_prompt = f"Rate the quality of this email on a scale of 1 to 10: \n\n{generated_email}"
-    score_response = generator(score_prompt, max_length=30, num_return_sequences=1)
-    score = score_response[0]['generated_text'].strip()
+        # Rate the email and suggest improvements
+        score_prompt = f"Rate the quality of this email on a scale of 1 to 10: \n\n{generated_email}"
+        score_response = generator(score_prompt, max_length=30, num_return_sequences=1)
+        score = int(score_response[0]['generated_text'].strip())
 
-    suggestion_prompt = f"Suggest specific improvements for this email: \n\n{generated_email}"
-    suggestion_response = generator(suggestion_prompt, max_length=50, num_return_sequences=1)
-    suggestion = suggestion_response[0]['generated_text'].strip()
+        suggestion_prompt = f"Suggest specific improvements for this email: \n\n{generated_email}"
+        suggestion_response = generator(suggestion_prompt, max_length=50, num_return_sequences=1)
+        suggestion = suggestion_response[0]['generated_text'].strip()
 
-    # Store the email, score, and suggestion
-    email_history.append({
-        "step": step_count,
-        "email": generated_email,
-        "score": score,
-        "suggestion": suggestion
-    })
+        # Store the email, score, and suggestion
+        email_history.append({
+            "step": step_count,
+            "email": generated_email,
+            "score": score,
+            "suggestion": suggestion
+        })
+
+        # Update previous draft and increment step count
+        previous_draft = generated_email
+        step_count += 1
 
     return generated_email, score, suggestion
 
@@ -54,16 +61,7 @@ def view_history():
 
 # Gradio interface function
 def email_agent(subject):
-    if not email_history:
-        step_count = 1
-        previous_draft = None
-    else:
-        step_count = email_history[-1]["step"] + 1
-        previous_draft = email_history[-1]["email"]
-        # Get the last suggestion for reflection
-        suggestion = email_history[-1]["suggestion"]
-
-    email, score, suggestion = generate_email(subject, previous_draft, step_count)
+    email, score, suggestion = generate_email(subject)
     history_display = view_history()
     
     return email, score, suggestion, history_display