From 92d6fca3231c8174d66514e20b0d6707442a46ab Mon Sep 17 00:00:00 2001 From: Xubin Ren Date: Tue, 14 Apr 2026 13:10:03 +0000 Subject: [PATCH] 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 --- nanobot/agent/loop.py | 7 +++++ nanobot/api/server.py | 41 +----------------------- nanobot/utils/document.py | 60 ++++++++++++++++++++++++++++++++++++ tests/test_api_attachment.py | 56 ++++++++++++++++++--------------- 4 files changed, 99 insertions(+), 65 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 1e9f1787..2f80cd94 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -35,6 +35,7 @@ from nanobot.command import CommandContext, CommandRouter, register_builtin_comm from nanobot.config.schema import AgentDefaults from nanobot.providers.base import LLMProvider from nanobot.session.manager import Session, SessionManager +from nanobot.utils.document import extract_documents from nanobot.utils.helpers import image_placeholder_text from nanobot.utils.helpers import truncate_text as truncate_text_fn from nanobot.utils.runtime import EMPTY_FINAL_RESPONSE_MESSAGE @@ -653,6 +654,12 @@ class AgentLoop: content=final_content or "Background task completed.", ) + # Extract document text from media at the processing boundary so all + # channels benefit without format-specific logic in ContextBuilder. + if msg.media: + new_content, image_only = extract_documents(msg.content, msg.media) + msg = dataclasses.replace(msg, content=new_content, media=image_only) + preview = msg.content[:80] + "..." if len(msg.content) > 80 else msg.content logger.info("Processing message from {}:{}: {}", msg.channel, msg.sender_id, preview) diff --git a/nanobot/api/server.py b/nanobot/api/server.py index fcba8b55..934879a3 100644 --- a/nanobot/api/server.py +++ b/nanobot/api/server.py @@ -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()) diff --git a/nanobot/utils/document.py b/nanobot/utils/document.py index 23e8eeee..e2aaeade 100644 --- a/nanobot/utils/document.py +++ b/nanobot/utils/document.py @@ -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 diff --git a/tests/test_api_attachment.py b/tests/test_api_attachment.py index ea8eed9d..6bd3676c 100644 --- a/tests/test_api_attachment.py +++ b/tests/test_api_attachment.py @@ -10,12 +10,12 @@ import pytest import pytest_asyncio from nanobot.api.server import ( - _extract_documents, _FileSizeExceeded, _parse_json_content, _save_base64_data_url, create_app, ) +from nanobot.utils.document import extract_documents try: from aiohttp.test_utils import TestClient, TestServer @@ -185,7 +185,7 @@ def test_parse_json_content_rejects_oversized_base64_file(tmp_path) -> None: @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") @pytest.mark.asyncio async def test_multipart_upload_saves_file(aiohttp_client, mock_agent, tmp_path) -> None: - """Multipart upload of non-image extracts text into content (not media).""" + """Multipart upload saves file to media dir and passes path to process_direct.""" import os original_cwd = os.getcwd() os.chdir(tmp_path) @@ -203,9 +203,8 @@ async def test_multipart_upload_saves_file(aiohttp_client, mock_agent, tmp_path) ) assert resp.status == 200 call_kwargs = mock_agent.process_direct.call_args.kwargs - assert "analyze this" in call_kwargs["content"] - # Non-image file text is extracted into content, not kept as media - assert not call_kwargs.get("media") + assert call_kwargs["content"] == "analyze this" + assert len(call_kwargs.get("media") or []) == 1 finally: os.chdir(original_cwd) @@ -373,7 +372,7 @@ async def test_json_base64_image_upload(aiohttp_client, mock_agent, tmp_path) -> # --------------------------------------------------------------------------- -# _extract_documents tests (API-layer document extraction) +# extract_documents tests (now in nanobot.utils.document) # --------------------------------------------------------------------------- def test_extract_documents_separates_images_from_docs(tmp_path) -> None: @@ -388,7 +387,7 @@ def test_extract_documents_separates_images_from_docs(tmp_path) -> None: docx_path = tmp_path / "report.docx" doc.save(docx_path) - text, image_paths = _extract_documents("summarize", [str(png), str(docx_path)]) + text, image_paths = extract_documents("summarize", [str(png), str(docx_path)]) assert len(image_paths) == 1 assert image_paths[0] == str(png) assert "Quarterly revenue" in text @@ -400,13 +399,13 @@ def test_extract_documents_skips_extraction_errors(tmp_path, monkeypatch) -> Non bad_file = tmp_path / "broken.docx" bad_file.write_text("not a docx", encoding="utf-8") - import nanobot.api.server as _srv + import nanobot.utils.document as _doc monkeypatch.setattr( - _srv, "extract_text", + _doc, "extract_text", lambda _path: "[error: failed to extract DOCX: boom]", ) - text, image_paths = _extract_documents("hello", [str(bad_file)]) + text, image_paths = extract_documents("hello", [str(bad_file)]) assert text == "hello" assert image_paths == [] @@ -415,23 +414,31 @@ def test_extract_documents_images_only(tmp_path) -> None: """When all files are images, text is unchanged and all paths kept.""" png = tmp_path / "a.png" png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100) - text, image_paths = _extract_documents("describe", [str(png)]) + text, image_paths = extract_documents("describe", [str(png)]) assert text == "describe" assert len(image_paths) == 1 +def test_extract_documents_skips_oversized_files(tmp_path) -> None: + """Files exceeding the size limit should be silently skipped.""" + big = tmp_path / "huge.txt" + big.write_bytes(b"x" * 200) + + text, image_paths = extract_documents("hello", [str(big)], max_file_size=100) + assert text == "hello" + assert image_paths == [] + + # --------------------------------------------------------------------------- -# DOCX end-to-end upload test (API layer now extracts text) +# DOCX upload test — API saves file, loop layer extracts text # --------------------------------------------------------------------------- @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") @pytest.mark.asyncio -async def test_docx_upload_extracted_and_sent(aiohttp_client, tmp_path) -> None: - """Uploaded DOCX text should be extracted at the API layer and - appended to the content string, not passed as media.""" - from docx import Document - - agent = _make_mock_agent("This report shows $5M revenue") +async def test_docx_upload_passes_media_path(aiohttp_client, tmp_path) -> None: + """Uploaded DOCX is saved to disk and its path passed as media. + (Text extraction happens later in AgentLoop._process_message.)""" + agent = _make_mock_agent("report summary") import os original_cwd = os.getcwd() os.chdir(tmp_path) @@ -440,25 +447,24 @@ async def test_docx_upload_extracted_and_sent(aiohttp_client, tmp_path) -> None: app = create_app(agent, model_name="m") client = await aiohttp_client(app) + from docx import Document doc = Document() - doc.add_heading("Q1 Report", level=1) doc.add_paragraph("Total revenue: $5,000,000") buf = BytesIO() doc.save(buf) - docx_bytes = buf.getvalue() import aiohttp data = aiohttp.FormData() data.add_field("message", "summarize the report") - data.add_field("files", docx_bytes, filename="report.docx", + data.add_field("files", buf.getvalue(), filename="report.docx", content_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document") resp = await client.post("/v1/chat/completions", data=data) assert resp.status == 200 call_kwargs = agent.process_direct.call_args.kwargs - # Document text should be extracted into content, not media - assert "Total revenue" in call_kwargs["content"] - # No media (docx is not an image) - assert not call_kwargs.get("media") + assert call_kwargs["content"] == "summarize the report" + media = call_kwargs.get("media", []) + assert len(media) == 1 + assert "report.docx" in media[0] finally: os.chdir(original_cwd)