hackless123/gmail_sys

1import random
2from .models import Email, Action, Observation, ActionType
3
4class GmailEnvironment:
5    def __init__(self):
6        self.max_steps = 5
7        self.reset()
8
9    def reset(self):
10        self.current_step = 0
11        # Synthetic email data
12        self.emails = [
13            Email(id="1", sender="hr@company.com", subject="Interview", body="Schedule a time.", is_spam_truth=False, is_priority_truth=True),
14            Email(id="2", sender="bot@spam.net", subject="WIN CASH", body="Click link now.", is_spam_truth=True, is_priority_truth=False),
15            Email(id="3", sender="mom@home.com", subject="Dinner?", body="Are you coming?", is_spam_truth=False, is_priority_truth=False),
16        ]
17        return self._get_obs()
18
19    def _get_obs(self):
20        # We create a clean version of emails for the agent by hiding the 'truth' fields
21        clean_inbox = []
22        for e in self.emails:
23            clean_email = Email(
24                id=e.id,
25                sender=e.sender,
26                subject=e.subject,
27                body=e.body
28                # We skip setting is_spam_truth and is_priority_truth here
29            )
30            clean_inbox.append(clean_email)
31            
32        return Observation(
33            inbox=clean_inbox, 
34            steps_taken=self.current_step, 
35            max_steps=self.max_steps
36        )
37
38    def step(self, action: Action):
39        reward = 0
40        target = next((e for e in self.emails if e.id == action.email_id), None)
41        
42        if target:
43            if action.action_type == ActionType.MARK_PRIORITY:
44                reward = 10 if target.is_priority_truth else -5
45            
46            elif action.action_type == ActionType.MARK_SPAM:
47                reward = 15 if target.is_spam_truth else -20
48            
49            elif action.action_type == ActionType.AUTO_REPLY:
50                if not target.is_spam_truth and not target.is_priority_truth:
51                    reward = 5
52                else:
53                    reward = -2
54            
55            elif action.action_type == ActionType.ARCHIVE:
56                reward = 0 # Neutral action
57
58        self.current_step += 1
59        done = self.current_step >= self.max_steps
60        return self._get_obs(), reward, done