agentic_implementation/schemas.py

# schemas.py

from pydantic import BaseModel, EmailStr
from typing import List, Literal, Optional, Union



class FetchEmailsParams(BaseModel):
    email: str
    query: str  # Changed from start_date/end_date to query for internal date extraction


class ShowEmailParams(BaseModel):
    message_id: str

class AnalyzeEmailsParams(BaseModel):
    emails: List[dict]

class DraftReplyParams(BaseModel):
    email: dict
    tone: Optional[Literal["formal", "informal"]] = "formal"

class SendReplyParams(BaseModel):
    message_id: str
    reply_body: str


ToolParams = Union[
    FetchEmailsParams,
    ShowEmailParams,
    AnalyzeEmailsParams,
    DraftReplyParams,
    SendReplyParams
]

class PlanStep(BaseModel):
    action: Literal[
        "fetch_emails",
        "show_email",
        "analyze_emails",
        "draft_reply",
        "send_reply",
        "done",
    ]
    parameters: Optional[ToolParams] = None
    
class Plan(BaseModel):
    plan: List[PlanStep]