-
Notifications
You must be signed in to change notification settings - Fork 3
/
main.py
347 lines (278 loc) · 15 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
import datetime
import imaplib
import smtplib
from email.mime.text import MIMEText
from email import policy
from email.parser import BytesParser
import json
import os
import time
import logging
from bs4 import BeautifulSoup
import re
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.sentiment import SentimentIntensityAnalyzer
import numpy as np
import email
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_groq import ChatGroq
import ollama
from config import USER, APP_PASSWORD, IMAP_URL, SMTP_URL, SMTP_PORT, SYSTEM_PROMPT
from groqtools import run_conversation
from getnow import get_current_time_formatted
from nltk.stem import WordNetLemmatizer
nltk.download('wordnet')
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('vader_lexicon')
now = get_current_time_formatted()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if GROQ_API_KEY is None:
raise ValueError("GROQ_API_KEY environment variable is not set.")
logging.basicConfig(filename='gmail_assistant.log', level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s')
class EmbeddingModel:
def __init__(self, json_file_path: str = 'conversation_memory.json', model_name: str = 'snowflake-arctic-embed'):
self.json_file_path = json_file_path
self.model_name = model_name
def _load_documents(self) -> list:
try:
with open(self.json_file_path, 'r') as file:
return json.load(file)
except FileNotFoundError:
return []
def get_embeddings(self) -> list:
documents = self._load_documents()
return [ollama.embeddings(model=self.model_name, prompt=doc['content'])['embedding'] for doc in documents]
class ConversationMemory:
def __init__(self, memory_file="conversation_memory.json"):
self.memory_file = memory_file
self.history = self._load_memory()
def save_context(self, role, content):
self.history.append({"role": role, "content": content})
self._save_memory()
def get_history(self):
return self.history
def _load_memory(self):
try:
with open(self.memory_file, "r") as file:
return json.load(file)
except FileNotFoundError:
return []
def _save_memory(self):
with open(self.memory_file, "w") as file:
json.dump(self.history, file, indent=2)
class ChatGroqFactory:
@staticmethod
def create_chat_groq(temperature=0.7, model_name="llama3-70b-8192"):
return ChatGroq(groq_api_key=GROQ_API_KEY, temperature=temperature, model_name=model_name)
def preprocess_text(text: str) -> str:
lemmatizer = WordNetLemmatizer()
tokens = word_tokenize(text)
stop_words = set(stopwords.words('english'))
lemmatized_tokens = [lemmatizer.lemmatize(token.lower()) for token in tokens if token.lower() not in stop_words]
return ' '.join(lemmatized_tokens)
def cosine_similarity(a, b):
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
class EmailHandler:
def __init__(self, embedding_model: EmbeddingModel):
self.embedding_model = embedding_model
self.sentiment_analyzer = SentimentIntensityAnalyzer()
self.imap_server = None
def retrieve_relevant_messages(self, user_prompt, memory):
embeddings = self.embedding_model.get_embeddings()
user_prompt_embedding = ollama.embeddings(model='snowflake-arctic-embed',
prompt=preprocess_text(user_prompt))['embedding']
return [
msg for i, msg in enumerate(memory.get_history())
if cosine_similarity(np.array(embeddings[i]), np.array(user_prompt_embedding)) > 0.8
]
def chat_with_groq(self, system_prompt, user_prompt, chat_instance=None, memory=None):
chat_instance = chat_instance or ChatGroqFactory.create_chat_groq()
memory = memory or ConversationMemory()
memory.save_context("user", user_prompt)
relevant_messages = self.retrieve_relevant_messages(user_prompt, memory)
history = relevant_messages + [{"role": "user", "content": user_prompt}]
messages = [SystemMessagePromptTemplate.from_template(system_prompt)] + [
HumanMessagePromptTemplate.from_template(msg["content"]) if msg["role"] == "user" else
SystemMessagePromptTemplate.from_template(msg["content"]) for msg in history
]
prompt = ChatPromptTemplate.from_messages(messages)
response = chat_instance.invoke(prompt.format_prompt())
memory.save_context("assistant", response.content)
return response
def get_email_body(self, msg):
if msg.is_multipart():
for part in msg.walk():
content_type = part.get_content_type()
if content_type == 'text/plain' and 'attachment' not in str(part.get('Content-Disposition')):
return part.get_payload(decode=True).decode('utf-8', errors='ignore')
elif content_type == 'text/html' and 'attachment' not in str(part.get('Content-Disposition')):
return self._parse_html_content(part.get_payload(decode=True).decode('utf-8', errors='ignore'))
else:
content_type = msg.get_content_type()
if content_type == 'text/plain':
return msg.get_payload(decode=True).decode('utf-8', errors='ignore')
elif content_type == 'text/html':
return self._parse_html_content(msg.get_payload(decode=True).decode('utf-8', errors='ignore'))
return ""
def _parse_html_content(self, html_content):
soup = BeautifulSoup(html_content, 'html.parser')
for element in soup(["script", "style"]):
element.decompose()
for link in soup.find_all('a'):
link.decompose()
text = soup.get_text()
return re.sub(r'\s+', ' ', text).strip()
def extract_sentences(self, text):
return sent_tokenize(text)
def clean_email_body(self, email_body):
lines = email_body.split("\n")
cleaned_lines = []
for line in lines:
if line.startswith("From:") or line.startswith("Sent:") or line.startswith("To:") or line.startswith(
"Subject:"):
break
cleaned_lines.append(line)
return "\n".join(cleaned_lines).strip()
def get_client_name(self, email_from):
if '<' in email_from:
client_name = email_from.split('<')[0].strip()
else:
client_name = email_from
return client_name
def process_email(self, mail_id, parser):
_, data = self.imap_server.fetch(mail_id, '(RFC822)')
for response_part in data:
if isinstance(response_part, tuple):
msg = parser.parsebytes(response_part[1])
email_from = msg['from']
email_subject = msg['subject'] or "No Subject"
email_body = self.get_email_body(msg)
print(f"Received email from: {email_from}")
print(f"Email subject: {email_subject}")
print(f"Email body: {email_body}")
skip_keywords = [
"no-reply", "do-not-reply", "noreply", "donotreply", "no_reply", "do_not_reply",
"newsletter", "news letter", "notifications", "notification", "account", "accounts"
]
if any(keyword in email_from.lower() or keyword in email_subject.lower() for keyword in skip_keywords):
logging.info(f"Skipping email from {email_from} with subject '{email_subject}'.")
print(f"Skipping email from {email_from} with subject '{email_subject}'.")
return
logging.info(f"Processing new email from {email_from} with subject '{email_subject}'.")
print(f"Processing new email from {email_from} with subject '{email_subject}'.")
cleaned_email_body = self.clean_email_body(email_body)
print(f"Cleaned email body: {cleaned_email_body}")
sentences = self.extract_sentences(cleaned_email_body)
client_name = self.get_client_name(email_from)
for sentence in sentences:
user_prompt = f" {now}\n\nAs Chaos, analyze this email sentence and determine if it requires a response. You will always respond to emails asking for help from you. You have RAG using embedding with NLP to find relevant info that's also attached.\n\nSender: {email_from}\n\nClient Name: {client_name}\n\nEmail Content:\n{sentence}\n\nDoes this email require a response? Respond with only YES or NO."
try:
response = self.chat_with_groq(SYSTEM_PROMPT, user_prompt)
print(f"Analysis result: {response.content}")
if "YES" in response.content.strip().upper():
user_prompt = f" {now}\n\nAs Chaos, respond to this email sentence. Use any tools necessary to formulate your response.\n\nSender: {email_from}\n\nClient Name: {client_name}\n\nEmail Content:\n{sentence}\n\nOnly include in your response what you are sending to the client."
print(f"Using tools to formulate a response...")
tool_response = run_conversation(user_prompt)
print(f"Tool response: {tool_response}")
response = self.chat_with_groq(SYSTEM_PROMPT, user_prompt + "\n" + tool_response)
print(f"Final response: {response.content}")
self.send_response_email(email_from, email_subject, response.content)
else:
user_prompt = f" {now}\n\nAs Chaos, analyze this email sentence and figure out what tools may be needed, if none, just say none or leave a note with the note tool.\n\nSender: {email_from}\n\nClient Name: {client_name}\n\nEmail Content:\n{sentence}\n\n"
tool_response = run_conversation(user_prompt)
logging.info(
f"No response required for email sentence from {email_from} with subject '{email_subject}'. Tool response: {tool_response}")
print(
f"No response required for email sentence from {email_from} with subject '{email_subject}'. Tool response: {tool_response}")
except Exception as e:
logging.error(f"An error occurred while processing the email sentence: {str(e)}")
print(f"An error occurred while processing the email sentence: {str(e)}")
continue
def save_draft_email(self, email_to, email_subject, email_body):
try:
message = MIMEText(email_body)
message['From'] = USER
message['To'] = email_to
message['Subject'] = email_subject
draft_folder = 'Drafts' # Specify the name of the draft folder in your email account
self.imap_server.append(draft_folder, '', imaplib.Time2Internaldate(time.time()), message.as_bytes())
logging.info(f"Draft email saved for: {email_to}")
print(f"Draft email saved for: {email_to}")
except Exception as e:
logging.error(f"An error occurred while saving the draft email: {str(e)}")
print(f"An error occurred while saving the draft email: {str(e)}")
def send_response_email(self, email_to, email_subject, response_content, send_now=False):
if send_now:
try:
smtp_server = smtplib.SMTP(SMTP_URL, SMTP_PORT)
smtp_server.starttls()
smtp_server.login(USER, APP_PASSWORD)
message = MIMEText(response_content)
message['From'] = USER
message['To'] = email_to
message['Subject'] = email_subject
smtp_server.sendmail(USER, [email_to], message.as_string())
logging.info(f"Email sent to: {email_to}")
print(f"Email sent to: {email_to}")
except Exception as e:
logging.error(f"An error occurred while sending the email: {str(e)}")
print(f"An error occurred while sending the email: {str(e)}")
finally:
smtp_server.quit()
else:
self.save_draft_email(email_to, email_subject, response_content)
def process_draft_emails(self):
try:
self.imap_server.select('Drafts') # Select the draft folder
_, data = self.imap_server.search(None, 'ALL')
mail_ids = data[0].split()
for mail_id in mail_ids:
_, data = self.imap_server.fetch(mail_id, '(RFC822)')
for response_part in data:
if isinstance(response_part, tuple):
msg = email.message_from_bytes(response_part[1])
email_to = msg['To']
email_subject = msg['Subject']
email_body = msg.get_payload()
self.send_response_email(email_to, email_subject, email_body, send_now=True)
self.imap_server.store(mail_id, '+FLAGS', '\\Deleted') # Mark the draft email as deleted
self.imap_server.expunge() # Permanently remove deleted draft emails
except Exception as e:
logging.error(f"An error occurred while processing draft emails: {str(e)}")
print(f"An error occurred while processing draft emails: {str(e)}")
def handle_emails(self):
parser = BytesParser(policy=policy.default)
while True:
try:
self.imap_server = imaplib.IMAP4_SSL(IMAP_URL)
self.imap_server.login(USER, APP_PASSWORD)
# Process incoming emails
self.imap_server.select('inbox')
_, data = self.imap_server.search(None, 'UNSEEN')
mail_ids = data[0].split()
if mail_ids:
for mail_id in mail_ids:
self.process_email(mail_id, parser)
else:
print("No new emails found.")
# Process draft emails
self.process_draft_emails()
except Exception as e:
logging.error(f"An error occurred while handling emails: {str(e)}")
print(f"An error occurred while handling emails: {str(e)}")
finally:
try:
self.imap_server.close()
self.imap_server.logout()
except Exception:
pass
print("Waiting for 10 seconds before checking for new emails...")
time.sleep(10) # Wait for 10 seconds before checking for new emails again
if __name__ == "__main__":
embedding_model = EmbeddingModel()
email_handler = EmailHandler(embedding_model)
email_handler.handle_emails()