From 5b31f6205246badde899b9f62542801dda896888 Mon Sep 17 00:00:00 2001 From: Rounak Bhatia Date: Mon, 9 Oct 2023 11:36:06 +0530 Subject: [PATCH] read email improved --- superagi/tools/email/read_email.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/superagi/tools/email/read_email.py b/superagi/tools/email/read_email.py index 702444f71..05bade982 100644 --- a/superagi/tools/email/read_email.py +++ b/superagi/tools/email/read_email.py @@ -7,6 +7,7 @@ from superagi.helper.imap_email import ImapEmail from superagi.helper.read_email import ReadEmail from superagi.helper.token_counter import TokenCounter +from superagi.lib.logger import logger from superagi.tools.base_tool import BaseTool @@ -15,6 +16,7 @@ class ReadEmailInput(BaseModel): page: int = Field(..., description="The index of the page result the function should resturn. Defaults to 0, the first page.") limit: int = Field(..., description="Number of emails to fetch in one cycle. Defaults to 5.") + received_from: str = Field(..., description="Email address of the specific sender from whom messages are to be fetched. Defaults to \"None\"") class ReadEmailTool(BaseTool): @@ -30,7 +32,7 @@ class ReadEmailTool(BaseTool): args_schema: Type[BaseModel] = ReadEmailInput description: str = "Read emails from an IMAP mailbox" - def _execute(self, imap_folder: str = "INBOX", page: int = 0, limit: int = 5) -> str: + def _execute(self, imap_folder: str = "INBOX", page: int = 0, limit: int = 5, received_from: str = "None") -> str: """ Execute the read email tool. @@ -53,13 +55,20 @@ def _execute(self, imap_folder: str = "INBOX", page: int = 0, limit: int = 5) -> status, messages = conn.select("INBOX") num_of_messages = int(messages[0]) messages = [] - for i in range(num_of_messages, num_of_messages - limit, -1): + count = 0 + for i in range(num_of_messages, 1, -1): # start from the latest email res, msg = conn.fetch(str(i), "(RFC822)") email_msg = {} for response in msg: self._process_message(email_msg, response) - messages.append(email_msg) - if TokenCounter.count_text_tokens(json.dumps(messages)) > self.max_token_limit: + + receiver_email = self.extract_email_from_header(email_msg["From"]) + if received_from == "None" or receiver_email == received_from: + messages.append(email_msg) + count += 1 + if count == limit: + break + if TokenCounter.count_text_tokens(json.dumps(messages)) > self.max_token_limit: break conn.logout() @@ -90,3 +99,8 @@ def _process_message(self, email_msg, response): body = msg.get_payload(decode=True).decode() if content_type == "text/plain": email_msg["Message Body"] = ReadEmail().clean_email_body(body) + + def extract_email_from_header(self, header_str): + start = header_str.find('<') + 1 + end = header_str.find('>', start) + return header_str[start:end]