inference.py

from typing import List, Literal
from pydantic import BaseModel

class EmailAction(BaseModel):
    action_type: Literal['send', 'reply', 'archive', 'delete']
    message: str = ""
    email_id: int = None

class EmailEnv:
    def __init__(self):
        self.reset()

    def reset(self):
        self.emails = [
            {"id": 1, "from": "alice@example.com", "subject": "Meeting Tomorrow", "label": "Work", "status": "Unread"},
            {"id": 2, "from": "bob@example.com", "subject": "Lunch Plans", "label": "Personal", "status": "Archived"},
        ]
        return {"emails": self.emails}

    def step(self, action: dict):
        if action["action_type"] == "send":
            new_id = len(self.emails) + 1
            self.emails.append({
                "id": new_id,
                "from": "me@example.com",
                "subject": "New Mail",
                "label": "Work",
                "status": "Sent"
            })
            return {"emails": self.emails, "reward": 1, "done": False}
        return {"emails": self.emails, "reward": 0, "done": False}

env = EmailEnv()
obs = None

def reset():
    global obs
    obs = env.reset()
    return obs

def step(action):
    global obs
    result = env.step(action)
    obs = result["emails"]
    return result

def act():
    return {"action_type": "send", "message": "Hello"}