refactor: centralize document extraction in AgentLoop._process_message

Move extract_documents() to nanobot.utils.document as a reusable helper
and call it once in AgentLoop._process_message, the single entry point
for all message processing (API + all channels).

This replaces the previous API-only _extract_documents() in server.py,
ensuring Telegram, Feishu, Slack, WeChat, and all other channels also
benefit from automatic document text extraction.

Adds a configurable max_file_size guard (default 50 MB) to skip
oversized files gracefully, preventing unbounded memory/CPU usage
from channel-downloaded attachments.

- server.py: removed _extract_documents and related imports
- document.py: added extract_documents() with size limit
- loop.py: calls extract_documents() at the top of _process_message
- Tests updated: 70 related tests pass

Made-with: Cursor
This commit is contained in:
Xubin Ren
2026-04-14 13:10:03 +00:00
parent 47f5795708
commit 92d6fca323
4 changed files with 99 additions and 65 deletions
+60
View File
@@ -1,9 +1,12 @@
"""Document text extraction utilities for nanobot."""
import mimetypes
from pathlib import Path
from loguru import logger
from nanobot.utils.helpers import detect_image_mime
try:
from pypdf import PdfReader
except ImportError:
@@ -204,3 +207,60 @@ def _is_text_extension(ext: str) -> bool:
".ini",
".cfg",
}
# ---------------------------------------------------------------------------
# High-level helper: split media into images + extracted document text
# ---------------------------------------------------------------------------
_MAX_EXTRACT_FILE_SIZE = 50 * 1024 * 1024 # 50 MB
def extract_documents(
text: str,
media_paths: list[str],
*,
max_file_size: int = _MAX_EXTRACT_FILE_SIZE,
) -> tuple[str, list[str]]:
"""Separate images from documents in *media_paths*.
Documents (PDF, DOCX, XLSX, PPTX, plain-text, …) have their text
extracted and appended to *text*. Only image paths are kept in the
returned list so that downstream layers only need to handle vision
blocks.
Files larger than *max_file_size* bytes are skipped with a warning
to avoid unbounded memory / CPU usage.
"""
image_paths: list[str] = []
doc_texts: list[str] = []
for path_str in media_paths:
p = Path(path_str)
if not p.is_file():
continue
try:
size = p.stat().st_size
except OSError:
continue
if size > max_file_size:
logger.warning(
"Skipping oversized file for extraction: {} ({:.1f} MB > {} MB limit)",
p.name, size / (1024 * 1024), max_file_size // (1024 * 1024),
)
continue
raw = p.read_bytes()
mime = detect_image_mime(raw) or mimetypes.guess_type(path_str)[0]
if mime and mime.startswith("image/"):
image_paths.append(path_str)
else:
extracted = extract_text(p)
if extracted and not extracted.startswith("[error:"):
doc_texts.append(f"[File: {p.name}]\n{extracted}")
if doc_texts:
text = text + "\n\n" + "\n\n".join(doc_texts)
return text, image_paths