@samitouri / QOSami_Agents_Zero-2 / commits / 3f6e55b6

added email_client helper.

TerminallyLazy committed Oct 16, 2025 at 18:16 UTC 3f6e55b65c949db1c9ab440e68c4d4f7f7446365
2 files changed +591 -1
python/helpers/email_client.py new
+587
@@ -0,0 +1,587 @@
1 +import asyncio
2 +import email
3 +import os
4 +import re
5 +import uuid
6 +from dataclasses import dataclass
7 +from email.header import decode_header
8 +from email.message import Message as EmailMessage
9 +from fnmatch import fnmatch
10 +from typing import Any, Dict, List, Optional, Tuple
11 +
12 +import html2text
13 +from bs4 import BeautifulSoup
14 +from imapclient import IMAPClient
15 +
16 +from python.helpers import files
17 +from python.helpers.errors import RepairableException, format_error
18 +from python.helpers.print_style import PrintStyle
19 +
20 +
21 +@dataclass
22 +class Message:
23 + """Email message representation with sender, subject, body, and attachments."""
24 + sender: str
25 + subject: str
26 + body: str
27 + attachments: List[str]
28 +
29 +
30 +class EmailClient:
31 + """
32 + Async email client for reading messages from IMAP and Exchange servers.
33 +
34 + """
35 +
36 + def __init__(
37 + self,
38 + account_type: str = "imap",
39 + server: str = "",
40 + port: int = 993,
41 + username: str = "",
42 + password: str = "",
43 + options: Optional[Dict[str, Any]] = None,
44 + ):
45 + """
46 + Initialize email client with connection parameters.
47 +
48 + Args:
49 + account_type: Type of account - "imap" or "exchange"
50 + server: Mail server address (e.g., "imap.gmail.com")
51 + port: Server port (default 993 for IMAP SSL)
52 + username: Email account username
53 + password: Email account password
54 + options: Optional configuration dict with keys:
55 + - ssl: Use SSL/TLS (default: True)
56 + - timeout: Connection timeout in seconds (default: 30)
57 + """
58 + self.account_type = account_type.lower()
59 + self.server = server
60 + self.port = port
61 + self.username = username
62 + self.password = password
63 + self.options = options or {}
64 +
65 + # Default options
66 + self.ssl = self.options.get("ssl", True)
67 + self.timeout = self.options.get("timeout", 30)
68 +
69 + self.client: Optional[IMAPClient] = None
70 + self.exchange_account = None
71 +
72 + async def connect(self) -> None:
73 + """Establish connection to email server."""
74 + try:
75 + if self.account_type == "imap":
76 + await self._connect_imap()
77 + elif self.account_type == "exchange":
78 + await self._connect_exchange()
79 + else:
80 + raise RepairableException(
81 + f"Unsupported account type: {self.account_type}. "
82 + "Supported types: 'imap', 'exchange'"
83 + )
84 + except Exception as e:
85 + err = format_error(e)
86 + PrintStyle.error(f"Failed to connect to email server: {err}")
87 + raise RepairableException(f"Email connection failed: {err}") from e
88 +
89 + async def _connect_imap(self) -> None:
90 + """Establish IMAP connection."""
91 + loop = asyncio.get_event_loop()
92 +
93 + def _sync_connect():
94 + client = IMAPClient(self.server, port=self.port, ssl=self.ssl, timeout=self.timeout)
95 + # Increase line length limit to handle large emails (default is 10000)
96 + # This fixes "line too long" errors for emails with large headers or embedded content
97 + client._imap._maxline = 100000
98 + client.login(self.username, self.password)
99 + return client
100 +
101 + self.client = await loop.run_in_executor(None, _sync_connect)
102 + PrintStyle.standard(f"Connected to IMAP server: {self.server}")
103 +
104 + async def _connect_exchange(self) -> None:
105 + """Establish Exchange connection."""
106 + try:
107 + from exchangelib import Account, Configuration, Credentials, DELEGATE
108 +
109 + loop = asyncio.get_event_loop()
110 +
111 + def _sync_connect():
112 + creds = Credentials(username=self.username, password=self.password)
113 + config = Configuration(server=self.server, credentials=creds)
114 + return Account(
115 + primary_smtp_address=self.username,
116 + config=config,
117 + autodiscover=False,
118 + access_type=DELEGATE
119 + )
120 +
121 + self.exchange_account = await loop.run_in_executor(None, _sync_connect)
122 + PrintStyle.standard(f"Connected to Exchange server: {self.server}")
123 + except ImportError as e:
124 + raise RepairableException(
125 + "exchangelib not installed. Install with: pip install exchangelib>=5.4.3"
126 + ) from e
127 +
128 + async def disconnect(self) -> None:
129 + """Clean up connection."""
130 + try:
131 + if self.client:
132 + loop = asyncio.get_event_loop()
133 + await loop.run_in_executor(None, self.client.logout)
134 + self.client = None
135 + PrintStyle.standard("Disconnected from IMAP server")
136 + elif self.exchange_account:
137 + self.exchange_account = None
138 + PrintStyle.standard("Disconnected from Exchange server")
139 + except Exception as e:
140 + PrintStyle.error(f"Error during disconnect: {format_error(e)}")
141 +
142 + async def read_messages(
143 + self,
144 + download_folder: str,
145 + filter: Optional[Dict[str, Any]] = None,
146 + ) -> List[Message]:
147 + """
148 + Read messages based on filter criteria.
149 +
150 + Args:
151 + download_folder: Folder to save attachments (relative to /a0/)
152 + filter: Filter criteria dict with keys:
153 + - unread: Boolean to filter unread messages (default: True)
154 + - sender: Sender pattern with wildcards (e.g., "*@company.com")
155 + - subject: Subject pattern with wildcards (e.g., "*invoice*")
156 + - since_date: Optional datetime for date filtering
157 +
158 + Returns:
159 + List of Message objects with attachments saved to download_folder
160 + """
161 + filter = filter or {}
162 +
163 + if self.account_type == "imap":
164 + return await self._fetch_imap_messages(download_folder, filter)
165 + elif self.account_type == "exchange":
166 + return await self._fetch_exchange_messages(download_folder, filter)
167 + else:
168 + raise RepairableException(f"Unsupported account type: {self.account_type}")
169 +
170 + async def _fetch_imap_messages(
171 + self,
172 + download_folder: str,
173 + filter: Dict[str, Any],
174 + ) -> List[Message]:
175 + """Fetch messages from IMAP server."""
176 + if not self.client:
177 + raise RepairableException("IMAP client not connected. Call connect() first.")
178 +
179 + loop = asyncio.get_event_loop()
180 + messages: List[Message] = []
181 +
182 + def _sync_fetch():
183 + # Select inbox
184 + self.client.select_folder("INBOX")
185 +
186 + # Build search criteria
187 + search_criteria = []
188 + if filter.get("unread", True):
189 + search_criteria.append("UNSEEN")
190 +
191 + if filter.get("since_date"):
192 + since_date = filter["since_date"]
193 + search_criteria.append(["SINCE", since_date])
194 +
195 + # Search for messages
196 + if not search_criteria:
197 + search_criteria = ["ALL"]
198 +
199 + message_ids = self.client.search(search_criteria)
200 + return message_ids
201 +
202 + message_ids = await loop.run_in_executor(None, _sync_fetch)
203 +
204 + if not message_ids:
205 + PrintStyle.hint("No messages found matching criteria")
206 + return messages
207 +
208 + PrintStyle.standard(f"Found {len(message_ids)} messages")
209 +
210 + # Fetch and process messages
211 + for msg_id in message_ids:
212 + try:
213 + msg = await self._fetch_and_parse_imap_message(msg_id, download_folder, filter)
214 + if msg:
215 + messages.append(msg)
216 + except Exception as e:
217 + PrintStyle.error(f"Error processing message {msg_id}: {format_error(e)}")
218 + continue
219 +
220 + return messages
221 +
222 + async def _fetch_and_parse_imap_message(
223 + self,
224 + msg_id: int,
225 + download_folder: str,
226 + filter: Dict[str, Any],
227 + ) -> Optional[Message]:
228 + """Fetch and parse a single IMAP message with retry logic for large messages."""
229 + loop = asyncio.get_event_loop()
230 +
231 + def _sync_fetch():
232 + try:
233 + # Try standard RFC822 fetch first
234 + return self.client.fetch([msg_id], ["RFC822"])[msg_id]
235 + except Exception as e:
236 + error_msg = str(e).lower()
237 + # If "line too long" error, try fetching in parts
238 + if "line too long" in error_msg or "fetch_failed" in error_msg:
239 + PrintStyle.warning(f"Message {msg_id} too large for standard fetch, trying alternative method")
240 + # Fetch headers and body separately to avoid line length issues
241 + try:
242 + envelope = self.client.fetch([msg_id], ["BODY.PEEK[]"])[msg_id]
243 + return envelope
244 + except Exception as e2:
245 + PrintStyle.error(f"Alternative fetch also failed for message {msg_id}: {format_error(e2)}")
246 + raise
247 + raise
248 +
249 + try:
250 + raw_msg = await loop.run_in_executor(None, _sync_fetch)
251 +
252 + # Extract email data from response
253 + if b"RFC822" in raw_msg:
254 + email_data = raw_msg[b"RFC822"]
255 + elif b"BODY[]" in raw_msg:
256 + email_data = raw_msg[b"BODY[]"]
257 + else:
258 + PrintStyle.error(f"Unexpected response format for message {msg_id}")
259 + return None
260 +
261 + email_msg = email.message_from_bytes(email_data)
262 +
263 + # Apply sender filter
264 + sender = self._decode_header(email_msg.get("From", ""))
265 + if filter.get("sender") and not fnmatch(sender, filter["sender"]):
266 + return None
267 +
268 + # Apply subject filter
269 + subject = self._decode_header(email_msg.get("Subject", ""))
270 + if filter.get("subject") and not fnmatch(subject, filter["subject"]):
271 + return None
272 +
273 + # Parse message
274 + return await self._parse_message(email_msg, download_folder)
275 +
276 + except Exception as e:
277 + PrintStyle.error(f"Failed to fetch/parse message {msg_id}: {format_error(e)}")
278 + return None
279 +
280 + async def _fetch_exchange_messages(
281 + self,
282 + download_folder: str,
283 + filter: Dict[str, Any],
284 + ) -> List[Message]:
285 + """Fetch messages from Exchange server."""
286 + if not self.exchange_account:
287 + raise RepairableException("Exchange account not connected. Call connect() first.")
288 +
289 + from exchangelib import Q
290 +
291 + loop = asyncio.get_event_loop()
292 + messages: List[Message] = []
293 +
294 + def _sync_fetch():
295 + # Build query
296 + query = None
297 + if filter.get("unread", True):
298 + query = Q(is_read=False)
299 +
300 + if filter.get("sender"):
301 + sender_pattern = filter["sender"].replace("*", "")
302 + sender_q = Q(sender__contains=sender_pattern)
303 + query = query & sender_q if query else sender_q
304 +
305 + if filter.get("subject"):
306 + subject_pattern = filter["subject"].replace("*", "")
307 + subject_q = Q(subject__contains=subject_pattern)
308 + query = query & subject_q if query else subject_q
309 +
310 + # Fetch messages from inbox
311 + inbox = self.exchange_account.inbox
312 + items = inbox.filter(query) if query else inbox.all()
313 + return list(items)
314 +
315 + exchange_messages = await loop.run_in_executor(None, _sync_fetch)
316 +
317 + PrintStyle.standard(f"Found {len(exchange_messages)} Exchange messages")
318 +
319 + # Process messages
320 + for ex_msg in exchange_messages:
321 + try:
322 + msg = await self._parse_exchange_message(ex_msg, download_folder)
323 + if msg:
324 + messages.append(msg)
325 + except Exception as e:
326 + PrintStyle.error(f"Error processing Exchange message: {format_error(e)}")
327 + continue
328 +
329 + return messages
330 +
331 + async def _parse_exchange_message(
332 + self,
333 + ex_msg,
334 + download_folder: str,
335 + ) -> Message:
336 + """Parse an Exchange message."""
337 + loop = asyncio.get_event_loop()
338 +
339 + def _get_body():
340 + return str(ex_msg.text_body or ex_msg.body or "")
341 +
342 + body = await loop.run_in_executor(None, _get_body)
343 +
344 + # Process HTML if present
345 + if ex_msg.body and str(ex_msg.body).strip().startswith("<"):
346 + body = self._html_to_text(str(ex_msg.body))
347 +
348 + # Save attachments
349 + attachment_paths = []
350 + if ex_msg.attachments:
351 + for attachment in ex_msg.attachments:
352 + if hasattr(attachment, "content"):
353 + path = await self._save_attachment_bytes(
354 + attachment.name,
355 + attachment.content,
356 + download_folder
357 + )
358 + attachment_paths.append(path)
359 +
360 + return Message(
361 + sender=str(ex_msg.sender.email_address) if ex_msg.sender else "",
362 + subject=str(ex_msg.subject or ""),
363 + body=body,
364 + attachments=attachment_paths
365 + )
366 +
367 + async def _parse_message(
368 + self,
369 + email_msg: EmailMessage,
370 + download_folder: str,
371 + ) -> Message:
372 + """
373 + Parse email message and extract content with inline attachments.
374 +
375 + Processes multipart messages, converts HTML to text, and maintains
376 + positional context for inline attachments.
377 + """
378 + sender = self._decode_header(email_msg.get("From", ""))
379 + subject = self._decode_header(email_msg.get("Subject", ""))
380 +
381 + # Extract body and attachments
382 + body = ""
383 + attachment_paths: List[str] = []
384 + cid_map: Dict[str, str] = {} # Map Content-ID to file paths
385 + body_parts: List[str] = [] # Track parts in order
386 +
387 + if email_msg.is_multipart():
388 + # Process parts in order to maintain attachment positions
389 + for part in email_msg.walk():
390 + content_type = part.get_content_type()
391 + content_disposition = str(part.get("Content-Disposition", ""))
392 +
393 + # Skip multipart containers
394 + if part.get_content_maintype() == "multipart":
395 + continue
396 +
397 + # Handle attachments
398 + if "attachment" in content_disposition or part.get("Content-ID"):
399 + filename = part.get_filename()
400 + if filename:
401 + filename = self._decode_header(filename)
402 + content = part.get_payload(decode=True)
403 + if content:
404 + path = await self._save_attachment_bytes(
405 + filename, content, download_folder
406 + )
407 + attachment_paths.append(path)
408 +
409 + # Map Content-ID for inline images
410 + cid = part.get("Content-ID")
411 + if cid:
412 + cid = cid.strip("<>")
413 + cid_map[cid] = path
414 +
415 + # Add positional marker for non-cid attachments
416 + # (cid attachments are positioned via HTML references)
417 + if not cid and body_parts:
418 + body_parts.append(f"\n[file://{path}]\n")
419 +
420 + # Handle body text
421 + elif content_type == "text/plain":
422 + if not body: # Use first text/plain as primary body
423 + charset = part.get_content_charset() or "utf-8"
424 + body = part.get_payload(decode=True).decode(charset, errors="ignore")
425 + body_parts.append(body)
426 +
427 + elif content_type == "text/html":
428 + if not body: # Use first text/html as primary body if no text/plain
429 + charset = part.get_content_charset() or "utf-8"
430 + html_content = part.get_payload(decode=True).decode(charset, errors="ignore")
431 + body = self._html_to_text(html_content, cid_map)
432 + body_parts.append(body)
433 +
434 + # Combine body parts if we built them up
435 + if len(body_parts) > 1:
436 + body = "".join(body_parts)
437 + else:
438 + # Single part message
439 + content_type = email_msg.get_content_type()
440 + charset = email_msg.get_content_charset() or "utf-8"
441 + content = email_msg.get_payload(decode=True)
442 + if content:
443 + if content_type == "text/html":
444 + body = self._html_to_text(content.decode(charset, errors="ignore"), cid_map)
445 + else:
446 + body = content.decode(charset, errors="ignore")
447 +
448 + return Message(
449 + sender=sender,
450 + subject=subject,
451 + body=body,
452 + attachments=attachment_paths
453 + )
454 +
455 + def _html_to_text(self, html_content: str, cid_map: Optional[Dict[str, str]] = None) -> str:
456 + """
457 + Convert HTML to plain text with inline attachment references.
458 +
459 + Replaces inline images with [file:///a0/...] markers to maintain
460 + positional context.
461 + """
462 + cid_map = cid_map or {}
463 +
464 + # Replace cid: references with file paths before conversion
465 + if cid_map:
466 + soup = BeautifulSoup(html_content, "html.parser")
467 + for img in soup.find_all("img"):
468 + src = img.get("src", "")
469 + if src.startswith("cid:"):
470 + cid = src[4:] # Remove "cid:" prefix
471 + if cid in cid_map:
472 + # Replace with file path marker
473 + file_marker = f"[file://{cid_map[cid]}]"
474 + img.replace_with(soup.new_string(file_marker))
475 + html_content = str(soup)
476 +
477 + # Convert HTML to text
478 + h = html2text.HTML2Text()
479 + h.ignore_links = False
480 + h.ignore_images = False
481 + h.ignore_emphasis = False
482 + h.body_width = 0 # Don't wrap lines
483 +
484 + text = h.handle(html_content)
485 +
486 + # Clean up extra whitespace
487 + text = re.sub(r"\n{3,}", "\n\n", text) # Max 2 consecutive newlines
488 + text = text.strip()
489 +
490 + return text
491 +
492 + async def _save_attachment_bytes(
493 + self,
494 + filename: str,
495 + content: bytes,
496 + download_folder: str,
497 + ) -> str:
498 + """
499 + Save attachment to disk and return absolute path.
500 +
501 + Uses Agent Zero's file helpers for path management.
502 + """
503 + # Sanitize filename
504 + filename = files.safe_file_name(filename)
505 +
506 + # Generate unique filename if needed
507 + unique_id = uuid.uuid4().hex[:8]
508 + name, ext = os.path.splitext(filename)
509 + unique_filename = f"{name}_{unique_id}{ext}"
510 +
511 + # Build relative path and save
512 + relative_path = os.path.join(download_folder, unique_filename)
513 + files.write_file_bin(relative_path, content)
514 +
515 + # Return absolute path
516 + abs_path = files.get_abs_path(relative_path)
517 + return abs_path
518 +
519 + def _decode_header(self, header: str) -> str:
520 + """Decode email header handling various encodings."""
521 + if not header:
522 + return ""
523 +
524 + decoded_parts = []
525 + for part, encoding in decode_header(header):
526 + if isinstance(part, bytes):
527 + decoded_parts.append(part.decode(encoding or "utf-8", errors="ignore"))
528 + else:
529 + decoded_parts.append(str(part))
530 +
531 + return " ".join(decoded_parts)
532 +
533 +
534 +async def read_messages(
535 + account_type: str = "imap",
536 + server: str = "",
537 + port: int = 993,
538 + username: str = "",
539 + password: str = "",
540 + download_folder: str = "tmp/email",
541 + options: Optional[Dict[str, Any]] = None,
542 + filter: Optional[Dict[str, Any]] = None,
543 +) -> List[Message]:
544 + """
545 + Convenience wrapper for reading email messages.
546 +
547 + Automatically handles connection and disconnection.
548 +
549 + Args:
550 + account_type: "imap" or "exchange"
551 + server: Mail server address
552 + port: Server port (default 993 for IMAP SSL)
553 + username: Email username
554 + password: Email password
555 + download_folder: Folder to save attachments (relative to /a0/)
556 + options: Optional configuration dict
557 + filter: Filter criteria dict
558 +
559 + Returns:
560 + List of Message objects
561 +
562 + Example:
563 + from python.helpers.email_client import read_messages
564 + messages = await read_messages(
565 + server="imap.gmail.com",
566 + port=993,
567 + username=secrets.get("EMAIL_USER"),
568 + password=secrets.get("EMAIL_PASSWORD"),
569 + download_folder="tmp/email/inbox",
570 + filter={"unread": True, "sender": "*@company.com"}
571 + )
572 + """
573 + client = EmailClient(
574 + account_type=account_type,
575 + server=server,
576 + port=port,
577 + username=username,
578 + password=password,
579 + options=options,
580 + )
581 +
582 + try:
583 + await client.connect()
584 + messages = await client.read_messages(download_folder, filter)
585 + return messages
586 + finally:
587 + await client.disconnect()
requirements.txt
+4 -1
@@ -4,7 +4,7 @@ browser-use==0.5.11
4 docker==7.1.0
5 duckduckgo-search==6.1.12
6 faiss-cpu==1.11.0
7 -fastmcp==2.3.4
7 +fastmcp==2.12.4
8 fasta2a==0.5.0
9 flask[async]==3.0.3
10 flask-basicauth==0.2.0
@@ -42,3 +42,6 @@ crontab==1.0.1
42 pathspec>=0.12.1
43 psutil>=7.0.0
44 soundfile==0.13.1
45 +imapclient>=3.0.1
46 +html2text>=2024.2.26
47 +beautifulsoup4>=4.12.3