diff --git a/llmail/__main__.py b/llmail/__main__.py index 2de4c9b..3238e07 100644 --- 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 index c75b194..29cdff4 100644 --- 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] \ No newline at end of file +__all__ = ["set_primary_logger", "logger", "args", bot_email] diff --git a/llmail/utils/cli_args.py b/llmail/utils/cli_args.py index c284daf..fbf4cbf 100644 --- 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,9 +103,7 @@ 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", @@ -117,12 +111,8 @@ def set_argparse(): 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 index bd99c9b..ee81524 100644 --- 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 index 5a7db80..ea282b0 100644 --- 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,9 +144,7 @@ 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 @@ -158,9 +152,7 @@ def get_thread_history( 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 index cade587..2a61dfb 100644 --- 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"))) -