Prompt refactoring V2 additional (#19)

* moved handle processing more modular and scaleable to custom handles

* poetry updates

* minor sub fixes to the overall prompt refactoring changes

* ruff linting changes

* did some redundant prompt removals and updated schedule handle prompts as suggested

* added custom exception for handle resolving issues

* fixed annottations issue in email agent

Author: kernelism

mxtoai/agents/email_agent.py

@@ -2,7 +2,7 @@
 import os
 import re
 from datetime import datetime
-from typing import Any
+from typing import Any, List, Dict
 
 from dotenv import load_dotenv
 
@@ -18,9 +18,15 @@
 )
 
 from mxtoai._logging import get_logger
-from mxtoai.prompts.base_prompts import create_attachment_processing_task, create_email_context, create_task_template
+from mxtoai.models import ProcessingInstructions
+from mxtoai.prompts.base_prompts import (
+    LIST_FORMATTING_REQUIREMENTS,
+    MARKDOWN_STYLE_GUIDE,
+    RESEARCH_GUIDELINES,
+    RESPONSE_GUIDELINES,
+)
 from mxtoai.routed_litellm_model import RoutedLiteLLMModel
-from mxtoai.schemas import EmailRequest
+from mxtoai.schemas import EmailRequest, EmailAttachment
 from mxtoai.scripts.report_formatter import ReportFormatter
 from mxtoai.scripts.visual_qa import azure_visualizer
 from mxtoai.tools.attachment_processing_tool import AttachmentProcessingTool
@@ -131,7 +137,7 @@ def __init__(
                     )
 
         # Define the list of tools available to the agent
-        self.available_tools: list[Tool] = [
+        self.available_tools: List[Tool] = [
             self.attachment_tool,
             self.schedule_tool,
             self.visit_webpage_tool,
@@ -168,7 +174,7 @@ def _init_agent(self):
 
         logger.debug("Agent initialized with routed model configuration")
 
-    def _get_required_actions(self, mode: str) -> list[str]:
+    def _get_required_actions(self, mode: str) -> List[str]:
         """Get list of required actions based on mode."""
         actions = []
         if mode in ["summary", "full"]:
@@ -202,7 +208,7 @@ def determine_mode_from_email(to_email: str) -> str:
 
         return mode_mapping.get(email_prefix, "full")
 
-    def _get_attachment_types(self, attachments: list[dict[str, Any]]) -> list[str]:
+    def _get_attachment_types(self, attachments: List[Dict[str, Any]]) -> List[str]:
         """Get list of attachment types."""
         types = []
         for att in attachments:
@@ -211,36 +217,74 @@ def _get_attachment_types(self, attachments: list[dict[str, Any]]) -> list[str]:
                 types.append(content_type)
         return types
 
-    def _create_task(self, email_request: EmailRequest, email_instructions: "EmailHandleInstructions") -> str:
+    def _create_task(self, email_request: EmailRequest, email_instructions: ProcessingInstructions) -> str:
         """Create a task description for the agent based on email handle instructions."""
-        # Create attachment details with explicit paths if needed
-        attachment_details = []
-        if email_instructions.process_attachments and email_request.attachments:
-            for att in email_request.attachments:
-                # Quote the file path to handle spaces correctly
-                quoted_path = f'"{att.path}"'
-                attachment_details.append(
-                    f"- {att.filename} (Type: {att.contentType}, Size: {att.size} bytes)\n"
-                    f"  EXACT FILE PATH: {quoted_path}"
-                )
-
-        # Create the email context section
-        email_context = create_email_context(email_request, attachment_details)
-
-        # Create attachment processing task if needed
-        attachment_task = create_attachment_processing_task(attachment_details) if attachment_details else ""
+        attachments = self._format_attachments(email_request.attachments) \
+            if email_instructions.process_attachments and email_request.attachments else []
 
-        # Create the complete task template
-        return create_task_template(
+        return self._create_task_template(
             handle=email_instructions.handle,
-            email_context=email_context,
+            email_context=self._create_email_context(email_request, attachments),
             handle_specific_template=email_instructions.task_template,
-            attachment_task=attachment_task,
+            attachment_task=self._create_attachment_task(attachments),
             deep_research_mandatory=email_instructions.deep_research_mandatory,
             output_template=email_instructions.output_template
         )
 
-    def _process_agent_result(self, final_answer_obj: Any, agent_steps: list) -> dict[str, Any]:
+    def _format_attachments(self, attachments: List[EmailAttachment]) -> List[str]:
+        """Format attachment details for inclusion in the task."""
+        return [
+            f"- {att.filename} (Type: {att.contentType}, Size: {att.size} bytes)\n"
+            f'  EXACT FILE PATH: "{att.path}"'
+            for att in attachments
+        ]
+    def _create_email_context(self, email_request: EmailRequest, attachment_details=None) -> str:
+        """Generate context information from the email request."""
+        recipients = ", ".join(email_request.recipients) if email_request.recipients else "N/A"
+        attachments_info = f"Available Attachments:\n{chr(10).join(attachment_details)}" if attachment_details else "No attachments provided."
+
+        return f"""Email Content:
+    Subject: {email_request.subject}
+    From: {email_request.from_email}
+    Email Date: {email_request.date}
+    Recipients: {recipients}
+    CC: {email_request.cc or "N/A"}
+    BCC: {email_request.bcc or "N/A"}
+    Body: {email_request.textContent or email_request.htmlContent or ""}
+
+    {attachments_info}
+    """
+
+    def _create_attachment_task(self, attachment_details: List[str]) -> str:
+        """Return instructions for processing attachments, if any."""
+        return f"Process these attachments:\n{chr(10).join(attachment_details)}" if attachment_details else ""
+
+    def _create_task_template(
+        self,
+        handle: str,
+        email_context: str,
+        handle_specific_template: str = "",
+        attachment_task: str = "",
+        deep_research_mandatory: bool = False,
+        output_template: str = "",
+    ) -> str:
+        """Combine all task components into the final task description."""
+        sections = [
+            f"Process this email according to the '{handle}' instruction type.\n",
+            email_context,
+            RESEARCH_GUIDELINES["mandatory"] if deep_research_mandatory else RESEARCH_GUIDELINES["optional"],
+            attachment_task,
+            handle_specific_template,
+            output_template,
+            RESPONSE_GUIDELINES,
+            MARKDOWN_STYLE_GUIDE,
+            LIST_FORMATTING_REQUIREMENTS
+        ]
+
+        return "\n\n".join(filter(None, sections))  # Filter out any empty strings
+
+
+    def _process_agent_result(self, final_answer_obj: Any, agent_steps: List) -> Dict[str, Any]:
         """
         Process the agent's result into our expected format, using the agent steps.
         Prioritizes direct output from the 'deep_research' tool if available.
@@ -604,7 +648,7 @@ def _process_agent_result(self, final_answer_obj: Any, agent_steps: list) -> dic
     def process_email(
         self,
         email_request: EmailRequest,
-        email_instructions: "EmailHandleInstructions",  # Type hint as string to avoid circular import
+        email_instructions: ProcessingInstructions,  # Type hint as string to avoid circular import
     ) -> dict[str, Any]:
         """
         Process an email using the agent based on the provided email handle instructions.

mxtoai/api.py

@@ -14,11 +14,11 @@
 from mxtoai._logging import get_logger
 from mxtoai.agents.email_agent import EmailAgent
 from mxtoai.config import ATTACHMENTS_DIR, SKIP_EMAIL_DELIVERY
+from mxtoai.dependencies import processing_instructions_resolver
 from mxtoai.email_sender import (
     generate_email_id,
     send_email_reply,
 )
-from mxtoai.handle_configuration import HANDLE_MAP
 from mxtoai.schemas import EmailAttachment, EmailRequest
 from mxtoai.tasks import process_email_task
 from mxtoai.validators import validate_api_key, validate_attachments, validate_email_handle, validate_email_whitelist
@@ -386,7 +386,7 @@ async def process_email(
             await file.seek(0)  # Reset file pointer for later use
 
         # Get handle configuration
-        email_instructions = HANDLE_MAP[handle]  # Safe to use direct access now
+        email_instructions = processing_instructions_resolver(handle)  # Safe to use direct access now
 
         # Log initial email details
         logger.info("Received new email request:")

mxtoai/dependencies.py

@@ -0,0 +1,4 @@
+from mxtoai.email_handles import DEFAULT_EMAIL_HANDLES
+from mxtoai.instruction_resolver import ProcessingInstructionsResolver
+
+processing_instructions_resolver = ProcessingInstructionsResolver(DEFAULT_EMAIL_HANDLES)

mxtoai/handle_configuration.py mxtoai/email_handles.pymxtoai/handle_configuration.py renamed to mxtoai/email_handles.py

@@ -1,28 +1,8 @@
-from typing import Optional
-
-from pydantic import BaseModel
-
+from mxtoai.models import ProcessingInstructions
 from mxtoai.prompts import output_prompts, template_prompts
 
-
-class EmailHandleInstructions(BaseModel):
-    handle: str
-    aliases: list[str]
-    process_attachments: bool
-    deep_research_mandatory: bool
-    rejection_message: Optional[str] = (
-        "This email handle is not supported. Please visit https://mxtoai.com/docs/email-handles to learn about supported email handles."
-    )
-    task_template: Optional[str] = None
-    requires_language_detection: bool = False  # Specifically for translate handle
-    requires_schedule_extraction: bool = False  # Specifically for schedule handle
-    target_model: Optional[str] = "gpt-4"  # Default to gpt-4, can be overridden per handle
-    output_template: Optional[str] = None  # Template for structuring the output
-
-
-# Define all email handle configurations
-EMAIL_HANDLES = [
-    EmailHandleInstructions(
+DEFAULT_EMAIL_HANDLES = [
+    ProcessingInstructions(
         handle="summarize",
         aliases=["summarise", "summary"],
         process_attachments=True,
@@ -31,7 +11,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.SUMMARIZE_TEMPLATE,
         output_template=output_prompts.SUMMARIZE_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="research",
         aliases=["deep-research"],
         process_attachments=True,
@@ -41,7 +21,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.RESEARCH_TEMPLATE,
         output_template=output_prompts.RESEARCH_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="simplify",
         aliases=["eli5", "explain"],
         process_attachments=True,
@@ -50,7 +30,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.SIMPLIFY_TEMPLATE,
         output_template=output_prompts.SIMPLIFY_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="ask",
         aliases=["custom", "agent", "assist", "assistant", "hi", "hello", "question"],
         process_attachments=True,
@@ -59,7 +39,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.ASK_TEMPLATE,
         output_template=output_prompts.ASK_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="fact-check",
         aliases=["factcheck", "verify"],
         process_attachments=True,
@@ -68,7 +48,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.FACT_TEMPLATE,
         output_template=output_prompts.FACT_CHECK_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="background-research",
         aliases=["background-check", "background"],
         process_attachments=True,
@@ -77,7 +57,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.BACKGROUND_RESEARCH_TEMPLATE,
         output_template=output_prompts.BACKGROUND_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="translate",
         aliases=["translation"],
         process_attachments=True,
@@ -86,7 +66,7 @@ class EmailHandleInstructions(BaseModel):
         task_template=template_prompts.TRANSLATE_TEMPLATE,
         output_template=output_prompts.TRANSLATION_OUTPUT_GUIDELINES,
     ),
-    EmailHandleInstructions(
+    ProcessingInstructions(
         handle="schedule",
         aliases=["schedule-action"],
         process_attachments=True,
@@ -97,10 +77,3 @@ class EmailHandleInstructions(BaseModel):
         output_template=output_prompts.SCHEDULE_OUTPUT_GUIDELINES,
     ),
 ]
-
-# Create a mapping of handles (including aliases) to their configurations
-HANDLE_MAP = {}
-for config in EMAIL_HANDLES:
-    HANDLE_MAP[config.handle] = config
-    for alias in config.aliases:
-        HANDLE_MAP[alias] = config

mxtoai/exceptions.py

@@ -0,0 +1,7 @@
+class UnspportedHandleException(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+class HandleAlreadyExistsException(Exception):
+    def __init__(self, message):
+        super().__init__(message)

mxtoai/instruction_resolver.py

@@ -0,0 +1,94 @@
+import mxtoai.exceptions as exceptions
+from mxtoai._logging import get_logger
+
+from .models import ProcessingInstructions
+
+logger = get_logger(__name__)
+
+class ProcessingInstructionsResolver:
+    """
+    Resolves processing instructions based on email handle names or aliases.
+
+    This class is responsible for:
+    - Registering predefined and custom email handle instructions.
+    - Resolving a given handle (or its alias) to its corresponding instructions.
+    - Preventing alias/handle conflicts unless explicitly allowed.
+    - Supporting runtime extensibility via dynamic handle registration.
+
+    Attributes:
+        handle_map (Dict[str, EmailHandleInstructions]): Maps handles and aliases to their instruction objects.
+
+    """
+
+    def __init__(self, default_instructions: list[ProcessingInstructions]):
+        """
+        Initialize the resolver with a list of default handle instructions.
+
+        Args:
+            default_instructions (List[EmailHandleInstructions]): List of predefined handle configurations.
+
+        """
+        self.handle_map: dict[str, ProcessingInstructions] = {}
+        self.register_handles(default_instructions)
+
+    def register_handles(self, instructions_list: list[ProcessingInstructions], overwrite: bool = False):
+        """
+        Registers a list of handle instructions (including aliases).
+
+        Args:
+            instructions_list (List[EmailHandleInstructions]): Handles and aliases to register.
+            overwrite (bool): If True, allows existing handles/aliases to be overwritten.
+
+        Raises:
+            ValueError: If a handle or alias is already registered and overwrite is False.
+
+        """
+        for instructions in instructions_list:
+            all_names = [instructions.handle, *instructions.aliases]
+            for name in all_names:
+                if name in self.handle_map and not overwrite:
+                    msg = f"Handle or alias '{name}' already registered. Use `overwrite=True` to replace."
+                    raise exceptions.HandleAlreadyExistsException(msg)
+                self.handle_map[name] = instructions
+
+    def add_custom_handle(self, custom_instruction: ProcessingInstructions, overwrite: bool = False):
+        """
+        Adds a single custom handle instruction.
+
+        Args:
+            custom_instruction (EmailHandleInstructions): The custom handle config to add.
+            overwrite (bool): If True, allows overwriting existing handles/aliases.
+
+        """
+        self.register_handles([custom_instruction], overwrite=overwrite)
+
+    def __call__(self, handle: str) -> ProcessingInstructions:
+        """
+        Resolves a handle or alias to its corresponding EmailHandleInstructions.
+
+        Args:
+            handle (str): The handle or alias name.
+
+        Returns:
+            EmailHandleInstructions: The matched instruction object.
+
+        Raises:
+            ValueError: If the handle is not registered.
+
+        """
+        if handle not in self.handle_map:
+            logger.debug("This email handle is not supported!")
+            raise exceptions.UnspportedHandleException(
+                "This email handle is not supported. Please visit https://mxtoai.com/docs/email-handles to learn about supported email handles."
+            )
+        return self.handle_map[handle]
+
+    def list_available_handles(self) -> list[str]:
+        """
+        Lists all unique handle names currently registered (excluding aliases).
+
+        Returns:
+            List[str]: A list of unique registered handle names.
+
+        """
+        return list({instr.handle for instr in self.handle_map.values()})