Initial commit

src/imap_sorter/exception.py

@@ -0,0 +1,11 @@
+class FolderDoesNotExistException(Exception):
+    """
+    Exception raised by the FakeImapConnection when a folder that was specified
+    does not exist.
+    """
+
+
+class MailDoesNotExistException(Exception):
+    """
+    Exception raised by the FakeImapConnection when a mail ID does not exist.
+    """

src/imap_sorter/hooks/logging.py

@@ -0,0 +1,70 @@
+import logging
+
+from pydantic_ai import (
+    ModelResponse,
+    ModelRequestContext,
+    RunContext,
+    UserPromptPart,
+    ToolReturnPart,
+    ThinkingPart,
+    ToolCallPart,
+    TextPart,
+    ToolDefinition,
+)
+from pydantic_ai.capabilities import Hooks, WrapModelRequestHandler, ValidatedToolArgs
+
+
+# Logger
+logger = logging.getLogger("imap_sorter.hooks.logging")
+
+
+class LoggingHooks:
+    """
+    Hook that logs requests, responses, and tool calls.
+    """
+
+    @property
+    def hooks(self) -> Hooks:
+        hooks = Hooks()
+        hooks.on.model_request(self.before_request)
+        hooks.on.before_tool_execute(self.before_tool_call)
+        return hooks
+
+    @staticmethod
+    async def before_request(
+        ctx: RunContext,
+        *,
+        request_context: ModelRequestContext,
+        handler: WrapModelRequestHandler,
+    ) -> ModelResponse:
+        for part in request_context.messages[-1].parts:
+            match part:
+                case UserPromptPart():
+                    logger.info("Prompt: %s", part.content)
+                case ToolReturnPart():
+                    pass
+                case _:
+                    logger.info("Request: %s", part)
+        resp = await handler(request_context)
+        for part in resp.parts:
+            match part:
+                case ThinkingPart():
+                    logger.info("Thinking: %s", part.provider_details["raw_content"][0])
+                case ToolCallPart():
+                    logger.debug("Calling tool %s", part.tool_name)
+                case TextPart():
+                    logger.info("Model: %s", part.content)
+                case _:
+                    logger.info("Model: %s", part)
+        return resp
+
+    @staticmethod
+    async def before_tool_call(
+        ctx: RunContext,
+        *,
+        call: ToolCallPart,
+        tool_def: ToolDefinition,
+        args: ValidatedToolArgs,
+    ) -> ValidatedToolArgs:
+        logger.info("Model used tool %s using %s", call.tool_name, call.args)
+        return args

src/imap_sorter/imap/connection.py

@@ -0,0 +1,141 @@
+import imaplib
+import logging
+import email
+import json
+from pathlib import Path
+
+from imap_sorter.exception import FolderDoesNotExistException, MailDoesNotExistException
+from imap_sorter.imap.models import EMailMetadata
+from imap_sorter.imap.utils import decode_rfc2231
+
+# Logger
+logger = logging.getLogger(__name__)
+
+
+def _parse_imap_folder(folder: bytes) -> tuple[bool, str]:
+    flags_raw, name_raw = folder.decode("utf-8").split('"/"')
+    flags = flags_raw.strip()[1:-1].replace("\\", "").split(" ")
+    return "HasNoChildren" in flags, name_raw.strip()
+
+
+class FakeImapConnection:
+    """
+    Class that mirrors a real IMAP connection
+    """
+
+    # List of folders on the IMAP server
+    _folders: list[str]
+
+    # The folder we are currently in
+    _current_folder: str
+
+    # Tracked info about emails
+    _emails: dict[str, EMailMetadata]
+
+    def __init__(self) -> None:
+        self._current_folder = "INBOX"
+        self._folders = []
+        self._emails = {}
+        self._operations = []
+
+    def sync_from_file(self, p: Path) -> None:
+        with p.open("r") as f:
+            data = json.load(f)
+
+        self._folders = data["folders"]
+        self._emails = {
+            mail["id"]: EMailMetadata(
+                id=mail["id"],
+                title=mail["title"],
+                folder=mail["folder"],
+            )
+            for mail in data["emails"]
+        }
+
+    @property
+    def current_folder(self) -> str:
+        return self._current_folder
+
+    def get_mails(self) -> list[EMailMetadata]:
+        return [
+            mail
+            for mail in self._emails.values()
+            if mail.folder == self._current_folder
+        ]
+
+    def get_folders(self) -> list[str]:
+        return self._folders
+
+    def create_folder(self, folder: str) -> None:
+        self._folders.append(folder)
+
+    def switch_folder(self, folder: str) -> None:
+        logger.info("Switching folder to %s", folder)
+        if folder not in self._folders:
+            raise FolderDoesNotExistException
+
+        self._current_folder = folder
+
+    def move_mail(self, mail_id: str, folder: str) -> None:
+        if folder not in self._folders:
+            raise FolderDoesNotExistException
+        if mail_id not in self._emails:
+            raise MailDoesNotExistException
+
+        self._emails[mail_id].folder = folder
+
+    def sync(self, host: str, port: int, username: str, password: str) -> None:
+        conn = imaplib.IMAP4_SSL(host=host, port=port)
+        conn.login(username, password)
+
+        # Find out what folders we have on the server
+        logger.info("Discovering folders...")
+        raw_folders = conn.list()[1]
+
+        self._folders = [
+            _parse_imap_folder(folder)[1] for folder in raw_folders if folder not in ()
+        ]
+        folders_to_exclude = ("Drafts", "Trash", "Sent", "Junk")
+        for f in folders_to_exclude:
+            if f in self._folders:
+                self._folders.remove(f)
+        logger.info(
+            "Found %d folders (excluding %d folders)",
+            len(self._folders),
+            len(folders_to_exclude),
+        )
+
+        # Discover emails in the folders
+        for folder in self._folders:
+            logger.debug("Selecting %s", folder)
+            conn.select(folder, readonly=True)
+            search_result = conn.search(None, "ALL")
+            msg_ids = [
+                mid
+                for mid in search_result[1][0].decode("utf-8").split(" ")
+                if mid != ""
+            ]
+            print(search_result)
+
+            logger.debug("Discovered %d emails in %s", len(msg_ids), folder)
+            for msg_id in msg_ids:
+                logger.debug("Fetching %s", msg_id)
+                result, data = conn.fetch(msg_id, "(RFC822)")
+                parsed = email.message_from_bytes(data[0][1])
+
+                self._emails[msg_id] = EMailMetadata(
+                    id=msg_id,
+                    title=decode_rfc2231(parsed["subject"]),
+                    folder=folder,
+                )
+
+        conn.logout()
+
+        with open("data.json", "w") as f:
+            json.dump(
+                {
+                    "emails": [mail.model_dump() for mail in self._emails.values()],
+                    "folders": self._folders,
+                },
+                f,
+            )

src/imap_sorter/imap/models.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+
+class EMailMetadata(BaseModel):
+    """
+    Metadata for an Email
+    """
+
+    title: str
+    id: str
+    folder: str

src/imap_sorter/imap/utils.py

@@ -0,0 +1,22 @@
+from email.header import decode_header
+
+
+def decode_rfc2231(string: str) -> str:
+    """
+    Decodes an RFC2231 (?) encoded string and decodes it.
+
+    Args:
+        string: The string to decode
+
+    Returns:
+        The decoded string.
+    """
+    parts = decode_header(string)
+    final_string = ""
+    for data, encoding in parts:
+        if isinstance(data, str):
+            final_string += data
+            continue
+
+        final_string += data.decode(encoding or "utf-8")
+    return final_string

src/imap_sorter/main.py

@@ -0,0 +1,119 @@
+import argparse
+import logging
+from pathlib import Path
+from typing import Any
+import json
+
+from pydantic_ai import RunContext, Agent
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.models.openai import OpenAIResponsesModel
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from imap_sorter.hooks.logging import LoggingHooks
+from imap_sorter.imap.connection import FakeImapConnection
+from imap_sorter.tools.imap import ImapTools
+
+
+class KnowsCurrentFolder(AbstractCapability[Any]):
+    """Tells the agent what the current folder is."""
+
+    _conn: FakeImapConnection
+
+    def __init__(self, conn: FakeImapConnection) -> None:
+        self._conn = conn
+        super().__init__()
+
+    def get_instructions(self):
+        def _get_folder(ctx: RunContext[Any]) -> str:
+            return f"You are currently in the IMAP folder '{self._conn.current_folder}'"
+
+        return _get_folder
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--username", type=str, help="IMAP username", required=True)
+    parser.add_argument("--password", type=str, help="IMAP password")
+    parser.add_argument("--host", "-H", type=str, help="IMAP host", required=True)
+    parser.add_argument("--port", "-P", type=int, help="IMAP port", required=True)
+    parser.add_argument(
+        "--data",
+        "-d",
+        type=Path,
+        help="Instead of querying the IMAP server, use this JSON file",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=Path,
+        help="Dump the agent's operations into this JSON file",
+    )
+    parser.add_argument(
+        "--url",
+        "-u",
+        help="URL to an OpenAI-compatible inference endpoint",
+        required=True,
+    )
+    parser.add_argument(
+        "--api-key",
+        help="API key to use with the OpenAI-compatible endpoint",
+        required=True,
+    )
+    parser.add_argument("--model", default="", help="Model name to use")
+    parsed = parser.parse_args()
+
+    logging.basicConfig(level=logging.DEBUG)
+    logging.getLogger("httpx").setLevel(logging.ERROR)
+    logging.getLogger("httpcore").setLevel(logging.ERROR)
+    logging.getLogger("openai").setLevel(logging.ERROR)
+    logging.getLogger("asyncio").setLevel(logging.ERROR)
+
+    conn = FakeImapConnection()
+    if parsed.data is not None:
+        conn.sync_from_file(parsed.data)
+    else:
+        conn.sync(
+            host=parsed.host,
+            port=parsed.port,
+            username=parsed.username,
+            password=parsed.password,
+        )
+
+    # Set up the agent.
+    hooks = LoggingHooks()
+    imap_tools = ImapTools(conn)
+    model = OpenAIResponsesModel(
+        parsed.model,
+        provider=OpenAIProvider(
+            base_url=parsed.url,
+            api_key=parsed.api_key,
+        ),
+    )
+    agent = Agent(
+        model,
+        capabilities=[
+            KnowsCurrentFolder(conn),
+            hooks.hooks,
+        ],
+        toolsets=[
+            imap_tools.toolset,
+        ],
+    )
+    agent.run_sync(
+        "Given the available tools, go through the email folders and decide how "
+        "the emails should be sorted into folders. You can move emails and create folders "
+        "as you think makes sense. Finish once every mail has been sorted into a folder that is not "
+        "'Inbox'. Keep in mind that you have to create folders before you move emails into them, if they "
+        "don't already exist! Try and avoid folder names that include hyphens (-) and underscores (_). If possible, "
+        "try and use single word descriptions.",
+    )
+
+    if parsed.output is not None:
+        with open("operations.json", "w") as f:
+            json.dump(
+                {
+                    "operations": imap_tools.dump_operations(),
+                },
+                f,
+                indent=4,
+            )

src/imap_sorter/scripts/preview.py

@@ -0,0 +1,66 @@
+import argparse
+import json
+from pathlib import Path
+from collections import defaultdict
+
+from imap_sorter.imap.models import EMailMetadata
+from imap_sorter.tools.operations import (
+    MailOperation,
+    FolderCreationOperation,
+    MailMovementOperation,
+)
+from imap_sorter.imap.utils import decode_rfc2231
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--operations",
+        type=Path,
+        help="Path to the operations file that the agent generated",
+        required=True,
+    )
+    parser.add_argument(
+        "--data", type=Path, help="Path to the data file the agent used", required=True
+    )
+    parsed = parser.parse_args()
+
+    # Read the emails
+    with parsed.data.open("r") as f:
+        data = json.load(f)
+        emails = [EMailMetadata(**m) for m in data["emails"]]
+
+    # Read the operations
+    with parsed.operations.open("r") as f:
+        data = json.load(f)
+        operations = [MailOperation.deserialize(o) for o in data["operations"]]
+
+    # Pre-process the emails
+    mail_folders: dict[str, str] = {}
+    for operation in operations:
+        match operation:
+            case FolderCreationOperation():
+                # Nothing to do here
+                pass
+            case MailMovementOperation():
+                mail_folders[operation.mail_id] = operation.destination
+
+    # Group them
+    grouped_emails: dict[str, list[EMailMetadata]] = defaultdict(list)
+    for email in emails:
+        email.folder = mail_folders[email.id]
+        grouped_emails[email.folder].append(email)
+
+    # Print them
+    for folder, folder_emails in grouped_emails.items():
+        print(f"::{folder}::")
+        for mail in folder_emails:
+            print(f"  {decode_rfc2231(mail.title).replace('\n', '')}")
+
+    # Print new folders
+    print("\n\nNew Folders:")
+    for operation in operations:
+        if not isinstance(operation, FolderCreationOperation):
+            continue
+
+        print(f"- {operation.name}")

src/imap_sorter/tools/imap.py

@@ -0,0 +1,89 @@
+from typing import Any
+
+from pydantic_ai import FunctionToolset
+
+from imap_sorter.imap.connection import FakeImapConnection
+from imap_sorter.exception import FolderDoesNotExistException, MailDoesNotExistException
+from imap_sorter.tools.operations import (
+    MailOperation,
+    FolderCreationOperation,
+    MailMovementOperation,
+)
+
+
+class ImapTools:
+    """
+    Provides the LLM with the tools to interact with the
+    simulated server.
+    """
+
+    # The fake IMAP connection.
+    _conn: FakeImapConnection
+
+    # Operations that the agent took.
+    _operations: list[MailOperation]
+
+    def __init__(self, connection: FakeImapConnection) -> None:
+        self._conn = connection
+        self._operations = []
+
+    @property
+    def toolset(self) -> FunctionToolset[Any]:
+        """
+        The IMAP toolset the agent can use.
+        """
+        return FunctionToolset(
+            tools=[
+                self.list_mails,
+                self.change_folder,
+                self.create_folder,
+                self.list_folders,
+                self.move_mail,
+            ],
+        )
+
+    def list_mails(self) -> list[tuple[str, str]]:
+        """Returns the emails in the current folder as a tuple of the ID of the email and the subject line."""
+        return [(mail.id, mail.title) for mail in self._conn.get_mails()]
+
+    def change_folder(self, folder: str) -> str:
+        """Changes the currently active folder."""
+        try:
+            self._conn.switch_folder(folder)
+            return "OK."
+        except FolderDoesNotExistException:
+            return f"Folder '{folder}' does not exist."
+
+    def create_folder(self, folder: str) -> None:
+        """Creates a new folder."""
+        self._conn.create_folder(folder)
+        self._operations.append(FolderCreationOperation(name=folder))
+
+    def list_folders(self) -> list[str]:
+        """Returns a list of known folders."""
+        return self._conn.get_folders()
+
+    def move_mail(self, mail_id: str, folder: str) -> str:
+        """Moves a mail into a new folder."""
+        try:
+            self._conn.move_mail(mail_id, folder)
+            self._operations.append(
+                MailMovementOperation(
+                    mail_id=mail_id,
+                    destination=folder,
+                )
+            )
+            return "OK."
+        except MailDoesNotExistException:
+            return f"The email with ID '{mail_id}' does not exist."
+        except FolderDoesNotExistException:
+            return f"The folder '{folder}' does not exist."
+
+    def dump_operations(self) -> list[dict]:
+        """
+        Turns the agent's operations into a list of JSON objects.
+
+        Returns:
+            A list of JSON objects that represents the modifications to the IMAP server.
+        """
+        return [o.model_dump() for o in self._operations]

src/imap_sorter/tools/operations.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class MailOperation(BaseModel):
+    """
+    Base-class for IMAP operations.
+    """
+
+    type: str
+
+    @staticmethod
+    def deserialize(data: dict[str, Any]) -> "MailOperation":
+        match data["type"]:
+            case "mail_movement":
+                return MailMovementOperation(**data)
+            case "folder_creation":
+                return FolderCreationOperation(**data)
+            case _:
+                raise ValueError(f"Invalid type '{data['type']}'")
+
+
+class MailMovementOperation(MailOperation):
+    """
+    Operation where the agent moved an email into a new folder.
+    """
+
+    mail_id: str
+    destination: str
+    type: str = "mail_movement"
+
+
+class FolderCreationOperation(MailOperation):
+    """
+    Operation where the agent created a new folder.
+    """
+
+    name: str
+    type: str = "folder_creation"