[Refactor] Update dependencies and loaders (#1062)
This commit is contained in:
Sidharth Mohanty
@@ -1,123 +1,142 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import quopri
|
||||
from email import message_from_bytes
|
||||
from email.utils import parsedate_to_datetime
|
||||
from textwrap import dedent
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
try:
|
||||
from llama_hub.gmail.base import GmailReader
|
||||
from google.auth.transport.requests import Request
|
||||
from google.oauth2.credentials import Credentials
|
||||
from google_auth_oauthlib.flow import InstalledAppFlow
|
||||
from googleapiclient.discovery import build
|
||||
except ImportError:
|
||||
raise ImportError("Gmail requires extra dependencies. Install with `pip install embedchain[gmail]`") from None
|
||||
raise ImportError(
|
||||
'Gmail requires extra dependencies. Install with `pip install --upgrade "embedchain[gmail]"`'
|
||||
) from None
|
||||
|
||||
from embedchain.loaders.base_loader import BaseLoader
|
||||
from embedchain.utils import clean_string
|
||||
|
||||
|
||||
def get_header(text: str, header: str) -> str:
|
||||
start_string_position = text.find(header)
|
||||
pos_start = text.find(":", start_string_position) + 1
|
||||
pos_end = text.find("\n", pos_start)
|
||||
header = text[pos_start:pos_end]
|
||||
return header.strip()
|
||||
class GmailReader:
|
||||
SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
|
||||
|
||||
def __init__(self, query: str, service=None, results_per_page: int = 10):
|
||||
self.query = query
|
||||
self.service = service or self._initialize_service()
|
||||
self.results_per_page = results_per_page
|
||||
|
||||
@staticmethod
|
||||
def _initialize_service():
|
||||
credentials = GmailReader._get_credentials()
|
||||
return build("gmail", "v1", credentials=credentials)
|
||||
|
||||
@staticmethod
|
||||
def _get_credentials():
|
||||
if not os.path.exists("credentials.json"):
|
||||
raise FileNotFoundError("Missing 'credentials.json'. Download it from your Google Developer account.")
|
||||
|
||||
creds = (
|
||||
Credentials.from_authorized_user_file("token.json", GmailReader.SCOPES)
|
||||
if os.path.exists("token.json")
|
||||
else None
|
||||
)
|
||||
|
||||
if not creds or not creds.valid:
|
||||
if creds and creds.expired and creds.refresh_token:
|
||||
creds.refresh(Request())
|
||||
else:
|
||||
flow = InstalledAppFlow.from_client_secrets_file("credentials.json", GmailReader.SCOPES)
|
||||
creds = flow.run_local_server(port=8080)
|
||||
with open("token.json", "w") as token:
|
||||
token.write(creds.to_json())
|
||||
return creds
|
||||
|
||||
def load_emails(self) -> List[Dict]:
|
||||
response = self.service.users().messages().list(userId="me", q=self.query).execute()
|
||||
messages = response.get("messages", [])
|
||||
|
||||
return [self._parse_email(self._get_email(message["id"])) for message in messages]
|
||||
|
||||
def _get_email(self, message_id: str):
|
||||
raw_message = self.service.users().messages().get(userId="me", id=message_id, format="raw").execute()
|
||||
return base64.urlsafe_b64decode(raw_message["raw"])
|
||||
|
||||
def _parse_email(self, raw_email) -> Dict:
|
||||
mime_msg = message_from_bytes(raw_email)
|
||||
return {
|
||||
"subject": self._get_header(mime_msg, "Subject"),
|
||||
"from": self._get_header(mime_msg, "From"),
|
||||
"to": self._get_header(mime_msg, "To"),
|
||||
"date": self._format_date(mime_msg),
|
||||
"body": self._get_body(mime_msg),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _get_header(mime_msg, header_name: str) -> str:
|
||||
return mime_msg.get(header_name, "")
|
||||
|
||||
@staticmethod
|
||||
def _format_date(mime_msg) -> Optional[str]:
|
||||
date_header = GmailReader._get_header(mime_msg, "Date")
|
||||
return parsedate_to_datetime(date_header).isoformat() if date_header else None
|
||||
|
||||
@staticmethod
|
||||
def _get_body(mime_msg) -> str:
|
||||
def decode_payload(part):
|
||||
charset = part.get_content_charset() or "utf-8"
|
||||
try:
|
||||
return part.get_payload(decode=True).decode(charset)
|
||||
except UnicodeDecodeError:
|
||||
return part.get_payload(decode=True).decode(charset, errors="replace")
|
||||
|
||||
if mime_msg.is_multipart():
|
||||
for part in mime_msg.walk():
|
||||
ctype = part.get_content_type()
|
||||
cdispo = str(part.get("Content-Disposition"))
|
||||
|
||||
if ctype == "text/plain" and "attachment" not in cdispo:
|
||||
return decode_payload(part)
|
||||
elif ctype == "text/html":
|
||||
return decode_payload(part)
|
||||
else:
|
||||
return decode_payload(mime_msg)
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
class GmailLoader(BaseLoader):
|
||||
def load_data(self, query):
|
||||
"""Load data from gmail."""
|
||||
if not os.path.isfile("credentials.json"):
|
||||
raise FileNotFoundError(
|
||||
"You must download the valid credentials file from your google \
|
||||
dev account. Refer this `https://cloud.google.com/docs/authentication/api-keys`"
|
||||
)
|
||||
|
||||
loader = GmailReader(query=query, service=None, results_per_page=20)
|
||||
documents = loader.load_data()
|
||||
logging.info(f"Gmail Loader: {len(documents)} mails found for query- {query}")
|
||||
def load_data(self, query: str):
|
||||
reader = GmailReader(query=query)
|
||||
emails = reader.load_emails()
|
||||
logging.info(f"Gmail Loader: {len(emails)} emails found for query '{query}'")
|
||||
|
||||
data = []
|
||||
data_contents = []
|
||||
logging.info(f"Gmail Loader: {len(documents)} mails found")
|
||||
for document in documents:
|
||||
original_size = len(document.text)
|
||||
for email in emails:
|
||||
content = self._process_email(email)
|
||||
data.append({"content": content, "meta_data": email})
|
||||
|
||||
snippet = document.metadata.get("snippet")
|
||||
meta_data = {
|
||||
"url": document.metadata.get("id"),
|
||||
"date": get_header(document.text, "Date"),
|
||||
"subject": get_header(document.text, "Subject"),
|
||||
"from": get_header(document.text, "From"),
|
||||
"to": get_header(document.text, "To"),
|
||||
"search_query": query,
|
||||
}
|
||||
return {"doc_id": self._generate_doc_id(query, data), "data": data}
|
||||
|
||||
# Decode
|
||||
decoded_bytes = quopri.decodestring(document.text)
|
||||
decoded_str = decoded_bytes.decode("utf-8", errors="replace")
|
||||
@staticmethod
|
||||
def _process_email(email: Dict) -> str:
|
||||
content = BeautifulSoup(email["body"], "html.parser").get_text()
|
||||
content = clean_string(content)
|
||||
return dedent(
|
||||
f"""
|
||||
Email from '{email['from']}' to '{email['to']}'
|
||||
Subject: {email['subject']}
|
||||
Date: {email['date']}
|
||||
Content: {content}
|
||||
"""
|
||||
)
|
||||
|
||||
# Slice
|
||||
mail_start = decoded_str.find("<!DOCTYPE")
|
||||
email_data = decoded_str[mail_start:]
|
||||
|
||||
# Web Page HTML Processing
|
||||
soup = BeautifulSoup(email_data, "html.parser")
|
||||
|
||||
tags_to_exclude = [
|
||||
"nav",
|
||||
"aside",
|
||||
"form",
|
||||
"header",
|
||||
"noscript",
|
||||
"svg",
|
||||
"canvas",
|
||||
"footer",
|
||||
"script",
|
||||
"style",
|
||||
]
|
||||
|
||||
for tag in soup(tags_to_exclude):
|
||||
tag.decompose()
|
||||
|
||||
ids_to_exclude = ["sidebar", "main-navigation", "menu-main-menu"]
|
||||
for id in ids_to_exclude:
|
||||
tags = soup.find_all(id=id)
|
||||
for tag in tags:
|
||||
tag.decompose()
|
||||
|
||||
classes_to_exclude = [
|
||||
"elementor-location-header",
|
||||
"navbar-header",
|
||||
"nav",
|
||||
"header-sidebar-wrapper",
|
||||
"blog-sidebar-wrapper",
|
||||
"related-posts",
|
||||
]
|
||||
|
||||
for class_name in classes_to_exclude:
|
||||
tags = soup.find_all(class_=class_name)
|
||||
for tag in tags:
|
||||
tag.decompose()
|
||||
|
||||
content = soup.get_text()
|
||||
content = clean_string(content)
|
||||
|
||||
cleaned_size = len(content)
|
||||
if original_size != 0:
|
||||
logging.info(
|
||||
f"[{id}] Cleaned page size: {cleaned_size} characters, down from {original_size} (shrunk: {original_size-cleaned_size} chars, {round((1-(cleaned_size/original_size)) * 100, 2)}%)" # noqa:E501
|
||||
)
|
||||
|
||||
result = f"""
|
||||
email from '{meta_data.get('from')}' to '{meta_data.get('to')}'
|
||||
subject: {meta_data.get('subject')}
|
||||
date: {meta_data.get('date')}
|
||||
preview: {snippet}
|
||||
content: f{content}
|
||||
"""
|
||||
data_content = dedent(result)
|
||||
data.append({"content": data_content, "meta_data": meta_data})
|
||||
data_contents.append(data_content)
|
||||
doc_id = hashlib.sha256((query + ", ".join(data_contents)).encode()).hexdigest()
|
||||
response_data = {"doc_id": doc_id, "data": data}
|
||||
return response_data
|
||||
@staticmethod
|
||||
def _generate_doc_id(query: str, data: List[Dict]) -> str:
|
||||
content_strings = [email["content"] for email in data]
|
||||
return hashlib.sha256((query + ", ".join(content_strings)).encode()).hexdigest()
|
||||
|
||||
Reference in New Issue
Block a user