Skip to content
This repository has been archived by the owner on Mar 1, 2024. It is now read-only.

Commit

Permalink
Update multion notebook
Browse files Browse the repository at this point in the history
  • Loading branch information
ajhofmann committed Aug 29, 2023
1 parent bf7cb47 commit b54332b
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 115 deletions.
31 changes: 11 additions & 20 deletions llama_hub/tools/gmail/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,15 +44,7 @@ def load_data(self) -> List[Document]:
"""Load emails from the user's account"""
self._cache_service()

messsages = self.search_messages()

results = []
for message in messsages:
text = message.pop("body")
extra_info = message
results.append(Document(text=text, extra_info=extra_info))

return results
return self.search_messages()

def _get_credentials(self) -> Any:
"""Get valid user credentials from storage.
Expand Down Expand Up @@ -87,10 +79,9 @@ def _get_credentials(self) -> Any:

return creds

def search_messages(self):
query = self.query

max_results = self.max_results
def search_messages(self, query: str, max_results: Optional[int] = None):
if not max_results:
max_results = self.max_results

self._cache_service()

Expand All @@ -102,17 +93,17 @@ def search_messages(self):
.get("messages", [])
)

result = []
results = []
try:
for message in messages:
message_data = self.get_message_data(message)
if not message_data:
continue
result.append(message_data)
text = message_data.pop("body")
extra_info = message_data
results.append(Document(text=text, extra_info=extra_info))
except Exception as e:
raise Exception("Can't get message data" + str(e))

return result
return results

def get_message_data(self, message):
message_id = message["id"]
Expand Down Expand Up @@ -161,14 +152,14 @@ def extract_message_body(self, message: dict):
from bs4 import BeautifulSoup

try:
body = base64.urlsafe_b64decode(message["raw"].encode("ASCII"))
body = base64.urlsafe_b64decode(message["raw"].encode("UTF-8"))
mime_msg = email.message_from_bytes(body)

# If the message body contains HTML, parse it with BeautifulSoup
if "text/html" in mime_msg:
soup = BeautifulSoup(body, "html.parser")
body = soup.get_text()
return body.decode("ascii")
return body.decode("UTF-8")
except Exception as e:
raise Exception("Can't parse message body" + str(e))

Expand Down
Loading

0 comments on commit b54332b

Please sign in to comment.