Update app.py

app.py

@@ -1,78 +1,72 @@
+import os
 import gradio as gr
-from transformers import pipeline
+import requests
 
-# Defensive model loading with error log
-try:
-    classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
-except Exception as e:
-    classifier = None
-    print(f"Error loading zero-shot classifier: {e}")
+# Get Groq API key from environment
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
-try:
-    generator = pipeline("text-generation", model="gpt2", max_length=150)
-except Exception as e:
-    generator = None
-    print(f"Error loading generator: {e}")
+# Groq endpoints (you may need to adjust according to the latest docs or your selected model)
+GROQ_COMPLETION_URL = "https://api.groq.com/openai/v1/chat/completions"  # for LLM chat/completion
 
-try:
-    summarizer = pipeline("summarization")
-except Exception as e:
-    summarizer = None
-    print(f"Error loading summarizer: {e}")
+# Utility: Groq API request for LLM tasks (chat/completion)
+def groq_completion(prompt, system_prompt=None):
+    headers = {
+        "Authorization": f"Bearer {GROQ_API_KEY}",
+        "Content-Type": "application/json"
+    }
+    body = {
+        "model": "mixtral-8x7b-32768",  # Or change to another supported Groq model
+        "messages": [
+            {"role": "system", "content": system_prompt or "You are a helpful assistant."},
+            {"role": "user", "content": prompt}
+        ],
+        "temperature": 0.3,
+        "max_tokens": 512
+    }
+    response = requests.post(GROQ_COMPLETION_URL, headers=headers, json=body)
+    response.raise_for_status()
+    return response.json()["choices"][0]["message"]["content"]
 
+# Use LLM prompts for multi-task extraction
 def email_assistant(email_text):
     try:
-        candidate_labels = ["Sales", "Support", "Spam", "Finance", "General"]
-        # Category & Confidence
-        if classifier:
-            classification = classifier(email_text, candidate_labels)
-            category = classification["labels"][0]
-            confidence = round(classification["scores"][0]*100, 2)
-        else:
-            category, confidence = "Error", "Error"
-        # Priority
+        # You can extract "Subject" and "From" programmatically, or just instruct the user to paste email in a template
+        prompt = f"""You will receive an email message. Analyze its metadata and text and extract:
+- Category (Sales, Support, Spam, Finance, General)
+- Confidence (as a percentage)
+- Priority (High/Normal/Low)
+- Suggested forwarding email or action
+- Professional draft response
+- Checklist of action items
+        
+Return results in JSON format.
+
+Email: {email_text}
+        """
+        result = groq_completion(prompt)
+        # Try to parse JSON for Gradio outputs, fallback to plain output
+        import json
         try:
-            # Defensive logic for type safety if confidence is not a float
-            priority = "High" if category in ["Sales", "Support"] and isinstance(confidence, float) and confidence > 80 else "Normal"
+            output = json.loads(result)
+            return [
+                output.get('Category', ''),
+                output.get('Confidence', ''),
+                output.get('Priority', ''),
+                output.get('Suggested Action', ''),
+                output.get('Draft Response', ''),
+                output.get('Action Items', ''),
+            ]
         except Exception:
-            priority = "Error"
-        # Suggested Action
-        forwarding_emails = {
-            "Sales": "sales@company.com",
-            "Support": "support@company.com",
-            "Finance": "finance@company.com",
-            "Spam": None,
-            "General": "info@company.com"
-        }
-        suggested_action = f"Forward to {forwarding_emails.get(category, 'Unknown')}" if forwarding_emails.get(category) else "No action suggested"
-        # Draft Response
-        response_prompt = f"Write a professional reply to this email:\n\n{email_text}\n\nResponse:"
-        if generator:
-            response = generator(response_prompt, max_length=150, num_return_sequences=1)[0]["generated_text"]
-        else:
-            response = "Error generating response."
-        # Action Items
-        if summarizer:
-            action_items = summarizer(email_text, max_length=50, min_length=10, do_sample=False)[0]["summary_text"]
-        else:
-            action_items = "Error extracting action items."
-        return [
-            category,
-            f"{confidence}%",
-            priority,
-            suggested_action,
-            response,
-            action_items
-        ]
+            # If not JSON, gracefully show whole model output in all fields
+            return [result]*6
     except Exception as e:
-        print(f"Error in main function: {e}")
-        return [
-            "Error", "Error", "Error", "Error", "Error", "Error"
-        ]
+        print(f"Error in Groq API call: {e}")
+        err = "Error"
+        return [err]*6
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Email Classification & Response Assistant")
-    email_input = gr.TextArea(label="Paste your email here")
+    gr.Markdown("# Email Classification & Response Assistant (Groq powered)")
+    email_input = gr.TextArea(label="Paste your email here (include Subject, From, Body if possible)")
     submit = gr.Button("Analyze Email")
     output_category = gr.Textbox(label="Category")
     output_confidence = gr.Textbox(label="Confidence")