Allow a custom model to be set with OPENAI_MODEL and --openai-model

.env.example

@@ -4,6 +4,9 @@ WATCH_INTERVAL=300
 OPENAI_API_KEY=""
 # At least when writing this, openrouter.ai has a free tier
 OPENAI_BASE_URL="https://openrouter.ai/api/v1"
+# The model to use for OpenAI
+# For openrouter.ai, you can check the available models at https://openrouter.ai/docs#models
+OPENAI_MODEL = "mistralai/mistral-7b-instruct:free"
 
 # Settings for Gmail
 IMAP_HOST=imap.gmail.com

llmail/__main__.py

@@ -67,14 +67,17 @@ def __repr__(self):
 
 args = None
 bot_email = None
+
 email_threads = {}
+
 SUBJECT = "autoreply password"
 
 
 def main():
     """Main entry point for the script."""
     global args
     global bot_email
+    global email_threads
     args = argparser.parse_args()
 
 
@@ -96,6 +99,8 @@ def main():
                 while True:
                     fetch_and_process_emails()
                     time.sleep(args.watch_interval)
+                    # **EMPTY THREADS**
+                    email_threads = {}
             else:
                 fetch_and_process_emails()
 
@@ -122,6 +127,7 @@ def fetch_and_process_emails():
             # messages = client.search([f"(OR SUBJECT \"{SUBJECT}\" SUBJECT \"Re: {SUBJECT}\")"])
             messages = client.search(["OR", "SUBJECT", SUBJECT, "SUBJECT", f"Re: {SUBJECT}"])
             for msg_id in messages:
+                # TODO: It seems this will throw a KeyError if an email is sent while this for loop is running. May have been fixed by emptying email_threads at the end of the while loop? This should be tested again to confirm
                 msg_data = client.fetch([msg_id], ["ENVELOPE", "BODY[]", "RFC822.HEADER"])
                 envelope = msg_data[msg_id][b"ENVELOPE"]
                 subject = envelope.subject.decode()
@@ -218,6 +224,7 @@ def fetch_and_process_emails():
                 message_id=message_id,
                 references_ids=references_ids,
                 openai=openai,
+                model=args.openai_model,
             )
 
         logger.debug(f"Keys in email_threads: {len(email_threads.keys())}")
@@ -395,7 +402,7 @@ def send_reply(
     message_id: str,
     references_ids: list[str],
     openai: OpenAI,
-    model="mistralai/mistral-7b-instruct:free",
+    model: str,
 ):
     """Send a reply to the email with the specified message ID."""
     # Set roles deletes the sender key so we need to store the sender before calling it

llmail/utils/cli_args.py

@@ -63,6 +63,11 @@ def set_argparse():
         help="OpenAI API endpoint",
         default=os.getenv("OPENAI_BASE_URL"),
     )
+    ai_api.add_argument(
+        "--openai-model",
+        help="Model to use for the LLM",
+        default=os.getenv("OPENAI_MODEL") if os.getenv("OPENAI_MODEL") else "mistralai/mistral-7b-instruct:free",
+    )
     # Email arguments
     email = argparser.add_argument_group("Email")
     email.add_argument(