read email improved

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]