style: Formatted code with Ruff

slashtechno · Jun 1, 2024 · 63c7f53 · 63c7f53
1 parent 63f90c8
commit 63c7f53
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 86 deletions.
diff --git a/llmail/__main__.py b/llmail/__main__.py
@@ -1,10 +1,8 @@
-
 import time
 from imapclient import IMAPClient
 from llmail.utils import logger, args, responding
 
 
-
 def main():
     """Main entry point for the script."""
     match args.subcommand:
@@ -16,11 +14,9 @@ def main():
                     print(folder[2])
         case None:
             logger.debug(args)
-            logger.info(f"Looking for emails that match the subject key \"{args.subject_key}\"")
+            logger.info(f'Looking for emails that match the subject key "{args.subject_key}"')
             if args.watch_interval:
-                logger.info(
-                    f"Watching for new emails every {args.watch_interval} seconds"
-                )
+                logger.info(f"Watching for new emails every {args.watch_interval} seconds")
                 while True:
                     responding.fetch_and_process_emails(
                         look_for_subject=args.subject_key,
@@ -38,7 +34,5 @@ def main():
                 )
 
 
-
-
 if __name__ == "__main__":
     main()
diff --git a/llmail/utils/__init__.py b/llmail/utils/__init__.py
@@ -2,4 +2,4 @@
 from llmail.utils.cli_args import args, bot_email
 
 # https://stackoverflow.com/a/31079085
-__all__ = ["set_primary_logger", "logger", "args", bot_email]
+__all__ = ["set_primary_logger", "logger", "args", bot_email]
diff --git a/llmail/utils/cli_args.py b/llmail/utils/cli_args.py
@@ -36,9 +36,7 @@ def set_argparse():
         title="Subcommands",
     )
     # Subcommand: list-folders
-    _ = subparsers.add_parser(
-        "list-folders", help="List all folders in the IMAP account and exit"
-    )
+    _ = subparsers.add_parser("list-folders", help="List all folders in the IMAP account and exit")
     # General arguments
     argparser.add_argument(
         "--log-level",
@@ -65,9 +63,7 @@ def set_argparse():
         "-w",
         help="Interval in seconds to check for new emails. If not set, will only check once.",
         type=int,
-        default=(
-            int(os.getenv("WATCH_INTERVAL")) if os.getenv("WATCH_INTERVAL") else None
-        ),
+        default=(int(os.getenv("WATCH_INTERVAL")) if os.getenv("WATCH_INTERVAL") else None),
     )
     # OpenAI-compatible API arguments
     ai_api = argparser.add_argument_group("OpenAI-compatible API")
@@ -107,22 +103,16 @@ def set_argparse():
         "--subject-key",
         "-s",
         help="Emails with this subject will be replied to",
-        default=(
-            os.getenv("SUBJECT_KEY") if os.getenv("SUBJECT_KEY") else "llmail autoreply"
-        ),
+        default=(os.getenv("SUBJECT_KEY") if os.getenv("SUBJECT_KEY") else "llmail autoreply"),
     )
     email.add_argument(
         "--alias",
         help="Name to use in the 'From' in addition to the email address",
         default=os.getenv("ALIAS") if os.getenv("ALIAS") else "LLMail",
     )
     imap = email.add_argument_group("IMAP")
-    imap.add_argument(
-        "--imap-host", help="IMAP server hostname", default=os.getenv("IMAP_HOST")
-    )
-    imap.add_argument(
-        "--imap-port", help="IMAP server port", default=os.getenv("IMAP_PORT")
-    )
+    imap.add_argument("--imap-host", help="IMAP server hostname", default=os.getenv("IMAP_HOST"))
+    imap.add_argument("--imap-port", help="IMAP server port", default=os.getenv("IMAP_PORT"))
     imap.add_argument(
         "--imap-username",
         help="IMAP server username",
@@ -134,12 +124,8 @@ def set_argparse():
         default=os.getenv("IMAP_PASSWORD"),
     )
     smtp = email.add_argument_group("SMTP")
-    smtp.add_argument(
-        "--smtp-host", help="SMTP server hostname", default=os.getenv("SMTP_HOST")
-    )
-    smtp.add_argument(
-        "--smtp-port", help="SMTP server port", default=os.getenv("SMTP_PORT")
-    )
+    smtp.add_argument("--smtp-host", help="SMTP server hostname", default=os.getenv("SMTP_HOST"))
+    smtp.add_argument("--smtp-port", help="SMTP server port", default=os.getenv("SMTP_PORT"))
     smtp.add_argument(
         "--smtp-username",
         help="SMTP server username",
@@ -153,9 +139,7 @@ def set_argparse():
     smtp.add_argument(
         "--message-id-domain",
         help="Domain to use for Message-ID header",
-        default=(
-            os.getenv("MESSAGE_ID_DOMAIN") if os.getenv("MESSAGE_ID_DOMAIN") else None
-        ),
+        default=(os.getenv("MESSAGE_ID_DOMAIN") if os.getenv("MESSAGE_ID_DOMAIN") else None),
     )
 
     check_required_args(
@@ -175,7 +159,7 @@ def set_argparse():
     )
     args = argparser.parse_args()
     # Setting bot_email instead of using imap_username directly in case support is needed for imap_username and bot_email being different
-    global bot_email 
+    global bot_email
     bot_email = args.imap_username
 
 

diff --git a/llmail/utils/responding.py b/llmail/utils/responding.py
@@ -10,13 +10,16 @@
 
 # Uses utils/__init__.py to import from utils/logging.py and utils/cli_args.py respectively
 from llmail.utils import logger, args, bot_email
+
 # Import files from utils/
 from llmail.utils import tracking
+
 # Import utilites from utils/utils.py
 from llmail.utils.utils import get_plain_email_content
 
 email_threads = {}
 
+
 def fetch_and_process_emails(
     look_for_subject: str,
     alias: str = None,
@@ -29,11 +32,7 @@ def fetch_and_process_emails(
         client.login(args.imap_username, args.imap_password)
 
         email_threads = {}
-        folders = (
-            args.folder
-            if args.folder
-            else [folder[2] for folder in client.list_folders()]
-        )
+        folders = args.folder if args.folder else [folder[2] for folder in client.list_folders()]
         # for folder in client.list_folders():
         # Disabling fetching from all folders due it not being inefficient
         # Instead, just fetch from INBOX and get the threads later
@@ -56,9 +55,7 @@ def fetch_and_process_emails(
             )
             for msg_id in messages:
                 # It seems this will throw a KeyError if an email is sent while this for loop is running. However, I think the real cause is when an email is deleted (via another client) while testing the code
-                msg_data = client.fetch(
-                    [msg_id], ["ENVELOPE", "BODY[]", "RFC822.HEADER"]
-                )
+                msg_data = client.fetch([msg_id], ["ENVELOPE", "BODY[]", "RFC822.HEADER"])
                 envelope = msg_data[msg_id][b"ENVELOPE"]
                 subject = envelope.subject.decode()
                 # Use regex to verify that the subject optionally starts with "Fwd: " or "Re: " and then the intended subject (nothing case-sensitive)
@@ -147,23 +144,15 @@ def fetch_and_process_emails(
                 message_id = email_thread.initial_email.message_id
                 msg_id = email_thread.initial_email.imap_id
                 references_ids = email_thread.initial_email.references
-            elif (
-                len(email_thread.replies) > 0
-                and email_thread.replies[-1].sender != bot_email
-            ):
+            elif len(email_thread.replies) > 0 and email_thread.replies[-1].sender != bot_email:
                 logger.debug(
                     f"Last email in thread for email {message_id} is from {email_thread.replies[-1].sender}"
                 )
                 message_id = email_thread.replies[-1].message_id
                 msg_id = email_thread.replies[-1].imap_id
                 references_ids = email_thread.replies[-1].references
-            elif (
-                len(email_thread.replies) > 0
-                and email_thread.replies[-1].sender == bot_email
-            ):
-                logger.debug(
-                    f"Last email in thread for email {message_id} is from the bot"
-                )
+            elif len(email_thread.replies) > 0 and email_thread.replies[-1].sender == bot_email:
+                logger.debug(f"Last email in thread for email {message_id} is from the bot")
                 continue
             else:
                 ValueError("Invalid email thread")
@@ -180,7 +169,7 @@ def fetch_and_process_emails(
                 model=args.openai_model,
             )
 
-        logger.info (f"Current number of email threads: {len(email_threads.keys())}")
+        logger.info(f"Current number of email threads: {len(email_threads.keys())}")
 
 
 def send_reply(
@@ -202,10 +191,14 @@ def send_reply(
     if system_prompt:
         thread.insert(0, {"role": "system", "content": system_prompt})
     references_ids.append(message_id)
-    generated_response = openai.chat.completions.create(
-        model=model,
-        messages=thread,
-    ).choices[0].message.content
+    generated_response = (
+        openai.chat.completions.create(
+            model=model,
+            messages=thread,
+        )
+        .choices[0]
+        .message.content
+    )
     logger.debug(f"Generated response: {generated_response}")
     yag = yagmail.SMTP(
         user={args.smtp_username: alias} if alias else args.smtp_username,
@@ -226,16 +219,14 @@ def send_reply(
         )
     except SSLError as e:
         if "WRONG_VERSION_NUMBER" in str(e):
-            logger.info(
-                "SSL error occurred. Trying to connect with starttls=True instead."
-            )
+            logger.info("SSL error occurred. Trying to connect with starttls=True instead.")
             yag = yagmail.SMTP(
                 user={args.smtp_username: alias} if alias else args.smtp_username,
                 password=args.smtp_password,
                 host=args.smtp_host,
                 port=int(args.smtp_port),
                 smtp_starttls=True,
-                smtp_ssl=False
+                smtp_ssl=False,
             )
             yag.send(
                 to=sender,
@@ -260,7 +251,6 @@ def send_reply(
     logger.debug(f"Thread history length: {len(thread)}")
 
 
-
 def set_roles(thread_history: list[dict]) -> list[dict]:
     """Change all email senders to roles (assistant or user)"""
     # Change email senders to roles

diff --git a/llmail/utils/tracking.py b/llmail/utils/tracking.py
@@ -43,15 +43,11 @@ def sort_replies(self):
         # However, this **significantly** reduces complexity so for now, it's fine
 
     def __repr__(self):
-        return (
-            f"EmailThread(initial_email={self.initial_email}, replies={self.replies})"
-        )
+        return f"EmailThread(initial_email={self.initial_email}, replies={self.replies})"
 
 
 class Email:
-    def __init__(
-        self, imap_id, message_id, subject, sender, timestamp, body, references
-    ):
+    def __init__(self, imap_id, message_id, subject, sender, timestamp, body, references):
         self.imap_id = imap_id
         self.message_id = message_id
         self.subject = subject
@@ -148,19 +144,15 @@ def get_thread_history(
                 {
                     "sender": get_sender(message)["email"],
                     "content": get_plain_email_content(message),
-                    "timestamp": make_tz_aware(
-                        parsedate_to_datetime(message.get("Date"))
-                    ),
+                    "timestamp": make_tz_aware(parsedate_to_datetime(message.get("Date"))),
                 }
             )
             message = prev_message
         # Sort the thread history by timestamp
         thread_history = sorted(thread_history, key=lambda x: x["timestamp"])
         return thread_history
     else:
-        raise TypeError(
-            "Invalid type for message. Must be an int, str, or EmailThread object."
-        )
+        raise TypeError("Invalid type for message. Must be an int, str, or EmailThread object.")
 
 
 def get_sender(message: Message) -> dict:
@@ -182,9 +174,7 @@ def get_top_level_email(client, msg_id, message_id=None):
 
     # Extract the References header and split it into individual message IDs
     references_header = headers.get("References", "")
-    references_ids = [
-        m_id.strip() for m_id in references_header.split() if m_id.strip()
-    ]
+    references_ids = [m_id.strip() for m_id in references_header.split() if m_id.strip()]
 
     # Extract the first message ID, which represents the top-level email in the thread
     # If it doesn't exist, use the current message ID. Not msg_id since msg_id is only for IMAP

diff --git a/llmail/utils/utils.py b/llmail/utils/utils.py
@@ -5,13 +5,13 @@
 
 from llmail.utils import logger
 
+
 def make_tz_aware(timestamp):
     # dt = parsedate_to_datetime(timestamp)
     dt = timestamp
     return dt.replace(tzinfo=timezone.utc) if dt.tzinfo is None else dt
 
 
-
 def get_plain_email_content(message: Message | str) -> str:
     """Get the content of the email message. If a  message object is provided, it will be parsed
     Otherwise, it is assumed that the content is already a string and will be converted to markdown.
@@ -25,14 +25,10 @@ def get_plain_email_content(message: Message | str) -> str:
                 try:
                     body = part.get_payload(decode=True)
                 except UnicodeDecodeError:
-                    logger.debug(
-                        "UnicodeDecodeError occurred. Trying to get payload as string."
-                    )
+                    logger.debug("UnicodeDecodeError occurred. Trying to get payload as string.")
                     body = str(part.get_payload())
                 if content_type == "text/plain":
-                    markdown = html2text.html2text(
-                        str(body.decode("unicode_escape"))
-                    ).strip()
+                    markdown = html2text.html2text(str(body.decode("unicode_escape"))).strip()
                     # logger.debug(f"Converted to markdown: {markdown}")
                     # if len(markdown) < 5:
                     #     logger.warning(
@@ -45,4 +41,3 @@ def get_plain_email_content(message: Message | str) -> str:
             # if len(body) < 5:
             #     logger.warning(f"Content is less than 5 characters | Content: {body}")
             return html2text.html2text(str(body.decode("unicode_escape")))
-