Update app.py

app.py

@@ -1,8 +1,6 @@
 import os
 import gradio as gr
-from langchain_community.llms import Groq
-from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
+from groq import Groq
 import pandas as pd
 from datetime import datetime
 from sendgrid import SendGridAPIClient
@@ -10,6 +8,9 @@ from sendgrid.helpers.mail import Mail
 import json
 from pathlib import Path
 from typing import Dict, List, Optional
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+from langchain.llms import Groq as LangchainGroq
 
 # Define the storage directory in the Hugging Face Space
 STORAGE_DIR = os.environ.get('STORAGE_DIR', 'data')
@@ -48,7 +49,6 @@ class UserProfile:
             with open(self.storage_path, 'r') as f:
                 data = json.load(f)
             
-            # Update existing profile or add new one
             profiles = data.get("profiles", [])
             profile_exists = False
             for i, profile in enumerate(profiles):
@@ -110,33 +110,31 @@ class EmailTemplate:
 
 class EmailGenie:
     def __init__(self):
-        self.llm = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+        groq_api_key = os.environ.get("GROQ_API_KEY")
+        self.client = Groq(api_key=groq_api_key)
         self.sendgrid_client = initialize_email_client()
         self.sender_email = os.environ.get("SENDER_EMAIL")
         self.user_profile = UserProfile()
         self.email_template = EmailTemplate()
         self.current_profile = None
-    
-    def set_current_profile(self, name: str):
-        self.current_profile = self.user_profile.get_profile(name)
-        return self.current_profile is not None
-    
-    def structure_info(self, info: str) -> str:
-        prompt = PromptTemplate(
+        
+        # Initialize LangChain components
+        self.llm = LangchainGroq(
+            groq_api_key=groq_api_key,
+            model_name="llama-3.1-8b-instant"
+        )
+        
+        # Create LLMChains with structured prompts
+        self.structure_info_prompt = PromptTemplate(
             input_variables=["info"],
             template="Please structure and summarize the following information about a person or company: {info}"
         )
-        chain = LLMChain(llm=self.llm, prompt=prompt)
-        return chain.run(info=info)
-
-    def generate_email(self, template_name: str, context: Dict) -> str:
-        if not self.current_profile:
-            return "Please set up your profile first."
-        
-        template = self.email_template.get_template(template_name)
-        structured_info = self.structure_info(context.get('recipient_info', ''))
+        self.structure_info_chain = LLMChain(
+            llm=self.llm,
+            prompt=self.structure_info_prompt
+        )
         
-        email_prompt = PromptTemplate(
+        self.email_generation_prompt = PromptTemplate(
             input_variables=["template", "sender_name", "sender_industry", "sender_background", "recipient", "subject", "structured_info"],
             template="""
             Using the following template and information, create a personalized email:
@@ -148,9 +146,26 @@ class EmailGenie:
             Recipient Information: {structured_info}
             """
         )
+        self.email_generation_chain = LLMChain(
+            llm=self.llm,
+            prompt=self.email_generation_prompt
+        )
+    
+    def set_current_profile(self, name: str):
+        self.current_profile = self.user_profile.get_profile(name)
+        return self.current_profile is not None
+    
+    def structure_info(self, info: str) -> str:
+        return self.structure_info_chain.run(info=info)
+
+    def generate_email(self, template_name: str, context: Dict) -> str:
+        if not self.current_profile:
+            return "Please set up your profile first."
         
-        chain = LLMChain(llm=self.llm, prompt=email_prompt)
-        return chain.run(
+        template = self.email_template.get_template(template_name)
+        structured_info = self.structure_info(context.get('recipient_info', ''))
+        
+        email_content = self.email_generation_chain.run(
             template=template,
             sender_name=self.current_profile['name'],
             sender_industry=self.current_profile['industry'],
@@ -159,6 +174,8 @@ class EmailGenie:
             subject=context['email_subject'],
             structured_info=structured_info
         )
+        
+        return email_content
 
     def preview_email(self, email_content: str) -> str:
         if not self.current_profile: