Update app.py

app.py

@@ -17,6 +17,20 @@ app.secret_key = os.environ.get("FLASK_SECRET_KEY")
 REDIRECT_URI = "https://huggingface.co/spaces/curiousgeorge1292/Custom_Profile_Email_Generator/oauth2callback"
 GMAIL_SCOPES = ["https://www.googleapis.com/auth/gmail.compose"]
 
+# Create a simple credential storage mechanism
+class CredentialStore:
+    def __init__(self):
+        self.credentials = {}
+
+    def set_credentials(self, user_email, creds):
+        self.credentials[user_email] = creds
+
+    def get_credentials(self, user_email):
+        return self.credentials.get(user_email)
+
+# Create a global instance
+credential_store = CredentialStore()
+
 # Retrieve JSON string from the environment variable
 client_secrets_json = os.environ.get("GMAIL_OAUTH_SECRET_JSON")
 
@@ -34,9 +48,18 @@ with open("temp_client_secrets.json", "w") as temp_file:
 flow = Flow.from_client_secrets_file(
         'temp_client_secrets.json',  # Path to the temporary JSON file
         scopes=['https://www.googleapis.com/auth/gmail.compose'],
-        redirect_uri="https://huggingface.co/spaces/curiousgeorge1292/Custom_Profile_Email_Generator/oauth2callback"
+        redirect_uri=REDIRECT_URI
 )
 
+def get_user_email_from_credentials(credentials):
+    try:
+        service = build('gmail', 'v1', credentials=credentials)
+        user_info = service.users().getProfile(userId='me').execute()
+        return user_info.get('emailAddress')
+    except Exception as e:
+        print(f"Error getting user email: {e}")
+        return None
+
 @app.route('/oauth2callback')
 def oauth2callback():
     flow = Flow.from_client_secrets_file(
@@ -48,7 +71,7 @@ def oauth2callback():
 
     # Save the credentials to the session
     credentials = flow.credentials
-    session['credentials'] = {
+    creds_dict = {
         'token': credentials.token,
         'refresh_token': credentials.refresh_token,
         'token_uri': credentials.token_uri,
@@ -56,26 +79,30 @@ def oauth2callback():
         'client_secret': credentials.client_secret,
         'scopes': credentials.scopes
     }
+    user_email = get_user_email_from_credentials(credentials)
+    if user_email:
+        credential_store.set_credentials(user_email, creds_dict)
+        
     return "Authorization successful. You can now save emails to Gmail drafts"
 
-@app.route('/authenticate_gmail')
-def authenticate_gmail():
+#@app.route('/authenticate_gmail')
+#def authenticate_gmail():
     # Gmail OAuth2 logic here
-    credentials = get_credentials_from_oauth()
-    if 'user_id' not in flask.session:
-        flask.session['user_id'] = str(uuid.uuid4())  # Generate a unique ID if not already present
-    user_id = flask.session.get('user_id')  # Get user-specific ID
-    redis_store.set(user_id, credentials)  # Save credentials to Redis (or a database)
+    #credentials = get_credentials_from_oauth()
+    #if 'user_id' not in flask.session:
+        #flask.session['user_id'] = str(uuid.uuid4())  # Generate a unique ID if not already present
+    #user_id = flask.session.get('user_id')  # Get user-specific ID
+    #redis_store.set(user_id, credentials)  # Save credentials to Redis (or a database)
     # After successful authentication:
-    session['credentials'] = credentials
-    return "Authentication successful!"
+    #session['credentials'] = credentials
+    #return "Authentication successful!"
 
-@app.route('/check_credentials')
-def check_credentials():
-    if 'credentials' in session:
-        return "Credentials are stored in session."
-    else:
-        return "No credentials found. Please authorize first."
+#@app.route('/check_credentials')
+#def check_credentials():
+    #if 'credentials' in session:
+        #return "Credentials are stored in session."
+    #else:
+        #return "No credentials found. Please authorize first."
 
 
 def save_to_gmail_drafts(credentials_info, subject, body, recipient):
@@ -296,17 +323,24 @@ def email_agent(credentials_info, name, email, prospect_name, linkedin_url, webs
         # Generate the email content
         email_content = generate_email(name, email, prospect_name, linkedin_url, website_url, context_url, word_count, email_purpose, interested_position, company_url, professional_title, personal_background)
 
-        # Save to Gmail drafts
-        try:            
-            save_to_gmail_drafts(credentials_info,
-                subject="Generated Email Subject", 
-                body=email_content, 
-                recipient=email  # Replace with the recipient's email
-                )
-            return ("<div style='color: green;'>Email saved to Gmail drafts successfully.</div>", email_content)
-        except Exception as e:
-            return (f"<div style='color: red;'>Error saving to Gmail drafts: {str(e)}</div>", email_content)
-        return email_content
+        # Get credentials from our store
+        credentials_info = credential_store.get_credentials(email)
+
+        if credentials_info:
+            # Save to Gmail drafts
+            try:            
+                save_to_gmail_drafts(credentials_info,
+                    subject="Generated Email Subject", 
+                    body=email_content, 
+                    recipient=email  # Replace with the recipient's email
+                    )
+                return ("<div style='color: green;'>Email saved to Gmail drafts successfully.</div>", email_content)
+            except Exception as e:
+                return (f"<div style='color: red;'>Error saving to Gmail drafts: {str(e)}</div>", email_content)
+
+        else:
+            return ("<div style='color: red;'>Please authenticate with Gmail first.</div>", 
+                    email_content)
         
     except Exception as e:
         # Handle errors during email generation
@@ -410,11 +444,7 @@ with gr.Blocks() as ui:
         return session.get('user_id', "default_user")
 
     email_button.click(
-        lambda name, email, prospect_name, linkedin_url, website_url, context_url, word_count, email_purpose, interested_position, company_url, professional_title, personal_background:
-        email_agent(
-            get_user_id(),  # Fetch and pass the user's ID dynamically
-            name, email, prospect_name, linkedin_url, website_url, context_url, word_count, email_purpose, interested_position, company_url, professional_title, personal_background
-        ),
+        email_agent,
         inputs=[name, email, prospect_name, linkedin_url, website_url, context_url, word_count, email_purpose, interested_position, company_url, professional_title, personal_background 
         ],
         outputs=[email_status, generated_email]