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
+1 -40
View File
@@ -19,8 +19,7 @@ from aiohttp import web
from loguru import logger
from nanobot.config.paths import get_media_dir
from nanobot.utils.document import extract_text
from nanobot.utils.helpers import detect_image_mime, safe_filename
from nanobot.utils.helpers import safe_filename
from nanobot.utils.runtime import EMPTY_FINAL_RESPONSE_MESSAGE
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB
@@ -162,40 +161,6 @@ async def _parse_multipart(request: web.Request) -> tuple[str, list[str], str |
return text, media_paths, session_id
# ---------------------------------------------------------------------------
# Pre-processing: extract document text at the API boundary
# ---------------------------------------------------------------------------
def _extract_documents(text: str, media_paths: list[str]) -> tuple[str, list[str]]:
"""Separate images from documents in *media_paths*.
Documents (PDF, DOCX, XLSX, PPTX, …) have their text extracted and
appended to *text*. Only image paths are kept in the returned list so
that downstream layers (ContextBuilder) only need to handle vision
blocks.
"""
image_paths: list[str] = []
doc_texts: list[str] = []
for path_str in media_paths:
p = Path(path_str)
if not p.is_file():
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
# ---------------------------------------------------------------------------
# Route handlers
# ---------------------------------------------------------------------------
@@ -232,10 +197,6 @@ async def handle_chat_completions(request: web.Request) -> web.Response:
logger.exception("Error parsing upload")
return _error_json(413, "File too large or invalid upload")
# Extract document text at the API boundary; only images stay in media.
if media_paths:
text, media_paths = _extract_documents(text, media_paths)
session_key = f"api:{session_id}" if session_id else API_SESSION_KEY
session_locks: dict[str, asyncio.Lock] = request.app["session_locks"]
session_lock = session_locks.setdefault(session_key, asyncio.Lock())