Mail Assistant

async def check_for_new_messages(self):
    """Check for new messages since last check"""
    try:           
        # Check for new messages since last history ID
        history = self.gmail_service.users().history().list(
            userId='me',
            startHistoryId=self.last_history_id,
            historyTypes=['messageAdded']
        ).execute()
        
        new_messages = []
        for history_record in history.get('history', []):
            for message_added in history_record.get('messagesAdded', []):
                message_id = message_added['message']['id']
                message_details = self.get_message_details(message_id)
                if message_details:
                    # Filter out draft messages
                    if not self.is_draft_message(message_details):
                        new_messages.append(message_details)
                    else:
                        print(f"📝 Skipping draft message: {message_details.get('subject', 'No Subject')}")
        
        if new_messages:
            print(f"\n🔔 Found {len(new_messages)} new message(s) at {datetime.now().strftime('%H:%M:%S')}")
            for message in new_messages:
                await self.generate_draft(message)
        
        # Update history ID
        if history.get('history'):
            self.last_history_id = history['history'][-1]['id']
        
    except HttpError as error:
        print(f"Error checking for new messages: {error}")


async def start_monitoring(self):
    """Start monitoring Gmail for new messages using polling"""   
    while True:
        await self.check_for_new_messages()
        await asyncio.sleep(10)  # Check new emails every 10 seconds

agent = Agent(
    name="email_response_generator", 
    model="openai/gpt-4o-mini",
    system_prompt="You are an assistant that helps draft professional email responses.",
    post_hook=self.save_draft,
)  

def _create_email_prompt(self, email):
    # Extract email information
    subject = email.get('subject', 'No Subject')
    sender = email.get('from', 'Unknown Sender')
    body = email.get('body', '')
    snippet = email.get('snippet', '')
    content = body if body else snippet
    
    prompt = f"""
You are an AI assistant that helps draft professional email responses. 

Please write a helpful and appropriate response to this email:

**Received Email:**
From: {sender}
Subject: {subject}
Content: {content}

**Instructions:**
- Write a professional response
- Do not assume anything that is not explicitly stated in the email
- Address any questions or requests appropriately
- Keep the tone friendly but business-appropriate
- If you need more information, ask clarifying questions
- End with a professional closing
- Keep the response concise but complete

**Your Response:**
"""
    return prompt.strip()


async def generate_draft(self, message):
    prompt = self._create_email_prompt(message)
    
    # Generate and create draft
    try:
        response_text = await self.agent(prompt=prompt, input_email=message).collect()
        if response_text:
            print("Generated Response:")
            print("-" * 40)
            print(response_text)
            print("-" * 40)
            
            print("Draft created successfully!")
            print("="*60)
        else:
            print("Failed to generate response")
    except Exception as e:
        print(f"Error generating/creating draft: {e}")
    
    print("="*60)

def save_draft(self):
    """Create a draft response to the original message"""
    span = get_run_context().current_span()
    original_message = span.input['input_email']
    response_text = span.output.collect_text()

    try:
        # Extract sender email from the original message
        from_header = original_message.get('from', '')
        if '<' in from_header and '>' in from_header:
            # Extract email from "Name <email@domain.com>" format
            sender_email = from_header.split('<')[1].split('>')[0].strip()
        else:
            # Assume the entire from header is the email
            sender_email = from_header.strip()
        
        # Create response subject
        original_subject = original_message.get('subject', '')
        if not original_subject.startswith('Re:'):
            response_subject = f"Re: {original_subject}"
        else:
            response_subject = original_subject
        
        # Create draft message
        draft_message = self._create_message(sender_email, response_subject, response_text)
        
        # Create the draft with thread ID to link it to the original conversation
        draft_body = {
            'message': {
                'raw': draft_message
            }
        }
        
        # Include thread ID if available to link the draft to the original conversation
        if original_message.get('thread_id'):
            draft_body['message']['threadId'] = original_message['thread_id']
        
        draft = self.gmail_service.users().drafts().create(
            userId='me',
            body=draft_body
        ).execute()
        
        print(f"Draft ID: {draft['id']}")
        print(f"To: {sender_email}")
        print(f"Subject: {response_subject}")
        return draft
        
    except Exception as e:
        print(f"Error creating draft: {e}")
        return None