knowpro.py: Email example (#1694)

Email import:
* Chunk large message bodies
* Gmail testing
* Test REPL improvements
* Bug fixes

Author: umeshma

tools/test_email.py

@@ -6,9 +6,16 @@
 import asyncio
 import sys
 import traceback
-from typing import Any, Iterable
+from typing import (
+    Any, 
+    Literal,
+    Iterable, 
+    Callable, 
+    Awaitable
+)
 from colorama import Fore
 from pathlib import Path
+import argparse
 
 import typechat
 
@@ -33,13 +40,30 @@
 from utool import print_result
 
 class EmailContext:
-    def __init__(self, db_path: str, conversation: EmailMemory) -> None:
-        self.db_path = db_path
+    def __init__(self, base_path: Path, db_name: str, conversation: EmailMemory) -> None:
+        self.base_path = base_path
+        self.db_path = base_path.joinpath(db_name)
         self.conversation = conversation
 
-    async def reset(self):
+    async def load_conversation(self, db_name: str, create_new:bool = False):
+        await self.conversation.settings.conversation_settings.storage_provider.close()
+        self.db_path = self.base_path.joinpath(db_name)
+        self.conversation = await load_or_create_email_index(str(self.db_path), create_new) 
+
+    # Delete the current conversation and re-create it
+    async def restart_conversation(self):
         await self.conversation.settings.conversation_settings.storage_provider.close()
-        self.conversation = await load_or_create_email_index(self.db_path, create_new=True) 
+        self.conversation = await load_or_create_email_index(str(self.db_path), create_new=True) 
+
+
+CommandHandler = Callable[[EmailContext, list[str]], Awaitable[None]]
+
+# Command decorator
+def command(parser: argparse.ArgumentParser):
+    def decorator(func: Callable):
+        func.parser = parser # type: ignore
+        return func
+    return decorator
 
 # Just simple test code
 # TODO : Once stable, move creation etc to utool.py
@@ -54,40 +78,46 @@ async def main():
 
     db_path = str(base_path.joinpath("pyEmails.db"))
     context = EmailContext(
-        db_path=db_path,
+        base_path,
+        "pyEmails.db",
         conversation=await load_or_create_email_index(db_path, create_new=False)
     )
     print(f"Using email memory at: {db_path}")
     await print_conversation_stats(context.conversation)
 
     # Command handlers
-    cmd_handlers = {
+    cmd_handlers: dict[str, CommandHandler] = {
+        "@exit": exit_app,
+        "@quit": exit_app,
         "@add_messages": add_messages,  # Add messages
+        "@parse_messages": parse_messages,
+        "@load_index": load_index,
         "@build_index": build_index, # Build index
         "@reset_index": reset_index, # Delete  index and start over
         "@search": search_index,   # Search index 
         "@answer": generate_answer # Question answer
     }
+    default_handler = generate_answer
     while True:
         line = input("✉>>").strip()
         if len(line) == 0:
             continue
-        elif line == "exit":
-            break        
-        args = shlex.split(line)
-        if len(args) < 1:
-            continue
         try:
+            args = shlex.split(line)
+            if len(args) < 1:
+                continue
             cmd = args[0].lower()
+            args.pop(0)
             if cmd == "@help":
-                print_commands(cmd_handlers.keys())
+                help(cmd_handlers, args)
             else:
                 cmd_handler = cmd_handlers.get(cmd)
+                if cmd_handler is None and not cmd.startswith("@"):
+                    cmd_handler = default_handler
                 if cmd_handler:
-                    args.pop(0)
                     await cmd_handler(context, args)
                 else:
-                    print_commands(cmd_handlers.keys())      
+                    print_commands(cmd_handlers)
         except Exception as e:
             print()
             print(Fore.RED, f"Error\n: {e}")
@@ -100,13 +130,26 @@ async def main():
 # ==
 
 # Adds messages. Takes a path either to a file or to a directory
+def _add_messages_def() -> argparse.ArgumentParser:
+    cmd = argparse.ArgumentParser(
+        description="Add messages to index"
+    )
+    cmd.add_argument(
+        "--path", 
+        default="", 
+        help="Path to an .eml file or to a directory with .eml files"
+    )
+    return cmd
+
+@command(_add_messages_def())
 async def add_messages(context: EmailContext, args: list[str]):
-    if len(args) < 1:
-        print_error("No path provided")
+    named_args = _add_messages_def().parse_args(args) 
+    if named_args.path is None:
+        print("No path provided")
         return
-    
+     
     # Get the path to the email file or directory of emails to ingest 
-    src_path = Path(args[0])
+    src_path = Path(named_args.path)
     emails: list[EmailMessage]
     if src_path.is_file():
         emails = [import_email_from_file(str(src_path))]
@@ -173,8 +216,79 @@ async def generate_answer(context: EmailContext, args:list[str]):
 
 async def reset_index(context: EmailContext, args: list[str]):
     print(f"Deleting {context.db_path}")
-    await context.reset()
+    await context.restart_conversation()
     await print_conversation_stats(context.conversation)
+
+
+def _load_index_def() -> argparse.ArgumentParser:
+    cmdDef = argparse.ArgumentParser(
+        description="Load index at given db path"
+    )
+    cmdDef.add_argument("--name", type=str, default="", help="Name of the index to load")
+    cmdDef.add_argument("--new", type=bool, default=False)
+    return cmdDef
+
+@command(_load_index_def())
+async def load_index(context: EmailContext, args: list[str]):
+    named_args = _load_index_def().parse_args(args)
+    
+    db_name: str = named_args.name
+    if (len(db_name) == 0):
+        return
+    
+    if not db_name.endswith(".db"):
+        db_name += ".db"
+    print(db_name)
+    await context.load_conversation(db_name, named_args.new)
+
+def _parse_messages_def() -> argparse.ArgumentParser:
+    cmdDef = argparse.ArgumentParser(
+        description="Parse messages in the given path"
+    )
+    cmdDef.add_argument("--path", type=str, default="")
+    cmdDef.add_argument("--verbose", type=bool, default=False)
+    return cmdDef
+
+@command(_parse_messages_def())
+async def parse_messages(context: EmailContext, args: list[str]):
+    named_args = _parse_messages_def().parse_args(args)
+    src_path = Path(named_args.path)
+    file_paths: list[str]
+    if src_path.is_file():
+        file_paths = [str(src_path)]
+    else:
+        file_paths = [str(file_path) for file_path in Path(src_path).iterdir() if file_path.is_file()]
+
+    print(f"Parsing {len(file_paths)} messages")
+    for file_path in file_paths:
+        try:
+            msg = import_email_from_file(file_path)
+            print(file_path)
+            if named_args.verbose:
+                print("####################")
+                print_email(msg)
+                print_knowledge(msg.get_knowledge())
+                print("####################")
+
+        except Exception as e:
+            print_error(file_path)
+            print_error(str(e))
+
+async def exit_app(context: EmailContext, args: list[str]):
+    print("Goodbye")
+    sys.exit(0)
+
+def help(handlers: dict[str, CommandHandler], args: list[str]):
+    if len(args) > 0:
+        cmd = handlers.get(args[0])
+        if cmd is not None:
+            print_help(cmd)
+            return
+        
+    print_commands(handlers)
+    print("@help <commandName> for details")
+
+
 #
 # Utilities
 #
@@ -203,11 +317,22 @@ def delete_sqlite_db(db_path: str):
         if os.path.exists(wal_path):
             os.remove(wal_path)
 
-
+#=========================
 #
 # Printing
 #
+#=========================
+
+def print_help(handler: CommandHandler):
+    if hasattr(handler, "parser"):
+        parser = argparse.ArgumentParser = handler.parser # type: ignore
+        print(parser.format_help())
+        print()
 
+def print_commands(commands: dict[str, CommandHandler]):
+    names = sorted(commands.keys())
+    print_list(Fore.GREEN, names, "COMMANDS", "ul")
+       
 def print_email(email: EmailMessage):
     print("From:", email.metadata.sender)
     print("To:", ", ".join(email.metadata.recipients))
@@ -234,15 +359,20 @@ def print_knowledge(knowledge: kplib.KnowledgeResponse):
     print()
     print(Fore.RESET)
 
-def print_commands(names: Iterable[str]):
-    print_list(Fore.GREEN, sorted(names), "Commands")
-
-def print_list(color, list: Iterable[Any], title: str):
+def print_list(color, list: Iterable[Any], title: str, type: Literal["plain", "ol", "ul"] = "plain"):
+    print(color)
     if title:
-        print(color + f"# {title}")
-        print()
-    for item in list:
-        print(color + " -", item)
+        print(f"# {title}\n")
+    if type == "plain":
+        for item in list:
+            print(item)
+    elif type == "ul":
+        for item in list:
+            print(f"- {item}")
+    elif type == "ol":
+        for i, item in enumerate(list):
+            print(f"{i + 1}. {item}")
+    print(Fore.RESET)
 
 def print_error(msg: str):
     print(Fore.RED + msg)

typeagent/emails/email_import.py

@@ -3,41 +3,42 @@
 
 import re
 from pathlib import Path
-from datetime import datetime
+from typing import Iterable
 
 from email import message_from_string
 from email.utils import parsedate_to_datetime
 from email.message import Message
 
 from .email_message import EmailMessage, EmailMessageMeta
 
-def import_emails_from_dir(dir_path: str) -> list[EmailMessage]:
+def import_emails_from_dir(dir_path: str, max_chunk_length: int = 4096) -> list[EmailMessage]:
     messages: list[EmailMessage] = []
     for file_path in Path(dir_path).iterdir():
-        messages.append(import_email_from_file(str(file_path.resolve())));
+        if file_path.is_file():
+            messages.append(import_email_from_file(str(file_path.resolve()), max_chunk_length))
     return messages
 
 # Imports an email file (.eml) as a list of EmailMessage objects
-def import_email_from_file(file_path: str) -> EmailMessage:
+def import_email_from_file(file_path: str, max_chunk_length: int = 4096) -> EmailMessage:
     email_string: str = ""
     with open(file_path, "r") as f:
         email_string = f.read()
 
-    return import_email_string(email_string)
+    return import_email_string(email_string, max_chunk_length)
 
 # Imports a single email MIME string and returns an EmailMessage object
-def import_email_string(email_string: str) -> EmailMessage:
+def import_email_string(email_string: str, max_chunk_length: int = 4096) -> EmailMessage:
     msg: Message = message_from_string(email_string)
-    email: EmailMessage = import_email_message(msg)
+    email: EmailMessage = import_email_message(msg, max_chunk_length)
     return email
 
-def import_forwarded_email_string(email_string: str) -> list[EmailMessage]:
+def import_forwarded_email_string(email_string: str, max_chunk_length: int = 4096) -> list[EmailMessage]:
     msg_parts = get_forwarded_email_parts(email_string)
-    return [import_email_string(part) for part in msg_parts if len(part) > 0]
+    return [import_email_string(part, max_chunk_length) for part in msg_parts if len(part) > 0]
 
 # Imports an email.message.Message object and returns an EmailMessage object
 # If the message is a reply, returns only the latest response. 
-def import_email_message(msg: Message) -> EmailMessage:
+def import_email_message(msg: Message, max_chunk_length: int) -> EmailMessage:
     # Extract metadata from
     email_meta = EmailMessageMeta(
         sender = msg.get("From", ""),
@@ -61,9 +62,10 @@ def import_email_message(msg: Message) -> EmailMessage:
     if email_meta.subject is not None:
         body = email_meta.subject + "\n\n" + body
 
+    body_chunks = _text_to_chunks(body, max_chunk_length)
     email: EmailMessage = EmailMessage(
         metadata=email_meta, 
-        text_chunks=[body], 
+        text_chunks=body_chunks, 
         timestamp=timestamp
     )
     return email
@@ -159,3 +161,35 @@ def _remove_empty(strings: list[str]) -> list[str]:
         if len(s) > 0:
             non_empty.append(s)
     return non_empty
+
+def _text_to_chunks(text: str, max_chunk_length: int) -> list[str]:
+    if len(text) < max_chunk_length:
+        return [text]
+    
+    paragraphs = _splitIntoParagraphs(text)
+    return list(_merge_chunks(paragraphs, "\n\n", max_chunk_length))
+
+def _splitIntoParagraphs(text: str) -> list[str]: 
+    return _remove_empty(re.split(r'\n{2,}', text))
+
+def _merge_chunks(chunks: Iterable[str], separator: str, max_chunk_length: int) -> Iterable[str]:
+    sep_length = len(separator)
+    cur_chunk:str = ""
+    for new_chunk in chunks:
+        cur_length = len(cur_chunk)
+        new_length = len(new_chunk)
+        if new_length > max_chunk_length:
+            # Truncate
+            new_chunk = new_chunk[0:max_chunk_length]
+            new_length = len(new_chunk)
+
+        if (cur_length + (new_length + sep_length) > max_chunk_length):
+            if cur_length > 0:
+                yield cur_chunk
+            cur_chunk = new_chunk
+        else:
+            cur_chunk += separator
+            cur_chunk += new_chunk
+
+    if (len(cur_chunk)) > 0:
+        yield cur_chunk

typeagent/emails/email_message.py

@@ -27,6 +27,14 @@ class EmailMessageMeta(IKnowledgeSource, IMessageMetadata):
     bcc: list[str] = Field(default_factory=list)
     subject: str | None = None
 
+    @property
+    def source(self) -> str | None:  # type: ignore[reportIncompatibleVariableOverride]
+        return self.sender
+
+    @property
+    def dest(self) -> str | list[str] | None:  # type: ignore[reportIncompatibleVariableOverride]
+        return self.recipients
+
     def get_knowledge(self) -> kplib.KnowledgeResponse:
         return kplib.KnowledgeResponse(
             entities=self.to_entities(),