Implement email retrieval functionality and enhance authentication flow in the API

app/controllers/mail.py

@@ -1,20 +1,19 @@
 """Module to search and list emails from Gmail."""
-import os
-import re
 import base64
 import hashlib
+import os
+import re
 from datetime import datetime, timedelta
 from venv import logger
-from ics import Calendar
 
-from langchain_core.documents import Document
+from ics import Calendar
 from langchain_community.document_loaders import (
+    CSVLoader,
     PyPDFLoader,
     UnstructuredExcelLoader,
-    CSVLoader,
     UnstructuredImageLoader,
 )
-
+from langchain_core.documents import Document
 from models.db import vectorstore
 
 SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
@@ -230,3 +229,71 @@ def collect(service, query=(datetime.today() - timedelta(days=10)).strftime("aft
         return f"{len(emails)} emails added to the collection."
     else:
         logger.info("No emails found after two weeks ago.")
+
+
+def get_emails(service, query=(datetime.today() - timedelta(days=10)).strftime("after:%Y/%m/%d"), max_results=10):
+    """
+    Retrieve a list of emails with subject, to, from, cc, and content.
+    
+    Args:
+        mailservice: Authenticated Gmail API service instance
+        max_results: Maximum number of emails to retrieve
+    
+    Returns:
+        List of dictionaries containing email details
+    """
+    try:
+        # List messages
+        response = service.users().messages().list(userId='me', q=query, maxResults=max_results).execute()
+        messages = response.get('messages', [])
+        email_list = []
+
+        if not messages:
+            return email_list
+
+        for message in messages:
+            # Get detailed message data
+            msg = service.users().messages().get(userId='me', id=message['id'], format='full').execute()
+            headers = msg['payload']['headers']
+
+            # Initialize email details
+            email_data = {
+                'subject': '',
+                'from': '',
+                'to': '',
+                'cc': '',
+                'content': '',
+                'snippet': msg['snippet'] if 'snippet' in msg else '',
+            }
+
+            # Extract headers
+            for header in headers:
+                name = header['name'].lower()
+                if name == 'subject':
+                    email_data['subject'] = header['value']
+                elif name == 'from':
+                    email_data['from'] = header['value']
+                elif name == 'to':
+                    email_data['to'] = header['value']
+                elif name == 'cc':
+                    email_data['cc'] = header['value']
+
+            # Extract content
+            if 'parts' in msg['payload']:
+                for part in msg['payload']['parts']:
+                    if part['mimeType'] == 'text/plain':
+                        email_data['content'] = base64.urlsafe_b64decode(part['body']['data']).decode('utf-8')
+                        break
+                    elif part['mimeType'] == 'text/html':
+                        email_data['content'] = base64.urlsafe_b64decode(part['body']['data']).decode('utf-8')
+                        break
+            elif 'data' in msg['payload']['body']:
+                email_data['content'] = base64.urlsafe_b64decode(msg['payload']['body']['data']).decode('utf-8')
+
+            email_list.append(email_data)
+
+        return email_list
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return []

app/router/auth.py

@@ -2,10 +2,10 @@
 import os
 import json
 import pickle
-from cv2 import log
 from fastapi import APIRouter, Request
 from fastapi.responses import JSONResponse
 from google_auth_oauthlib.flow import InstalledAppFlow
+from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import build
 
 router = APIRouter(tags=["auth"])
@@ -74,25 +74,19 @@ async def google_callback(state: str, code: str, scope: str, request: Request):
     flow.redirect_uri = REDIRECT_URI
     flow.fetch_token(code=code)
     credentials = flow.credentials
+    print(credentials.to_json())
     request.state.session["credential"] = json.loads(credentials.to_json())
-    # cred_dict = (request.state.session.get("credential"))
-    # cred = Credentials(
-    #     token=cred_dict["token"],
-    #     refresh_token=cred_dict["refresh_token"],
-    #     token_uri=cred_dict["token_uri"],
-    #     client_id=cred_dict["client_id"],
-    #     client_secret=cred_dict["client_secret"],
-    #     scopes=cred_dict["scopes"],
-    # )
-    # service = build("gmail", "v1", credentials=Credentials(
-    #     token=cred_dict["token"],
-    #     refresh_token=cred_dict["refresh_token"],
-    #     token_uri=cred_dict["token_uri"],
-    #     client_id=cred_dict["client_id"],
-    #     client_secret=cred_dict["client_secret"],
-    #     scopes=cred_dict["scopes"],
-    # ))
-    service = build("gmail", "v1", credentials=credentials)
+    cred_dict = (request.state.session.get("credential"))
+    service = build("gmail", "v1", credentials=Credentials(
+        token=cred_dict["token"],
+        # refresh_token=cred_dict["refresh_token"],
+        # token_uri=cred_dict["token_uri"],
+        # client_id=cred_dict["client_id"],
+        # client_secret=cred_dict["client_secret"],
+        # scopes=cred_dict["scopes"],
+    ))
+    print(service)
+    # service = build("gmail", "v1", credentials=credentials)
     profile = service.users().getProfile(userId="me").execute()
     with open(f"cache/{profile['emailAddress']}.pickle", "wb") as token:
         pickle.dump(credentials, token)

app/router/mail.py

@@ -46,16 +46,32 @@ def collect(query: MailReqData, request: Request):
         logger.error("Error collecting mail: %s", e)
         return JSONResponse(content={"error": str(e)}, status_code=500)
 
-# @router.get("")
-# def get():
-#     """
-#     Handles the chat POST request.
-
-#     Args:
-#         query (ReqData): The request data containing the query parameters.
-
-#     Returns:
-#         str: The generated response from the chat function.
-#     """
-#     # result = mail.get()
-#     return JSONResponse(content= result)
+@router.get("")
+def get(query: MailReqData, request: Request):
+    """
+    Handles the chat POST request.
+
+    Args:
+        query (ReqData): The request data containing the query parameters.
+
+    Returns:
+        str: The generated response from the chat function.
+    """
+    if os.path.exists(f"cache/{query.email}.pickle"):
+            with open(f"cache/{query.email}.pickle", "rb") as token:
+                credentials = pickle.load(token)
+    else:
+        cred_dict = request.state.session.get("credential")
+        credentials = Credentials(
+            token=cred_dict["token"],
+            refresh_token=cred_dict["refresh_token"],
+            token_uri=cred_dict["token_uri"],
+            client_id=cred_dict["client_id"],
+            client_secret=cred_dict["client_secret"],
+            scopes=cred_dict["scopes"],
+        )
+    mailservice = build("gmail", "v1", credentials=credentials)
+    print("query: ", query.query)
+    result = mail.get_emails(mailservice, query.query)
+    print(result)
+    return JSONResponse(content= result)