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

@@ -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)
 

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 safe_filename
-from nanobot.utils.helpers import detect_image_mime, 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())

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

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"]
+        assert call_kwargs["content"] == "analyze this"
-        # Non-image file text is extracted into content, not kept as media
+        assert len(call_kwargs.get("media") or []) == 1
-        assert not call_kwargs.get("media")
     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:
+async def test_docx_upload_passes_media_path(aiohttp_client, tmp_path) -> None:
-    """Uploaded DOCX text should be extracted at the API layer and
+    """Uploaded DOCX is saved to disk and its path passed as media.
-    appended to the content string, not passed as media."""
+    (Text extraction happens later in AgentLoop._process_message.)"""
-    from docx import Document
+    agent = _make_mock_agent("report summary")
-
-    agent = _make_mock_agent("This report shows $5M revenue")
     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 call_kwargs["content"] == "summarize the report"
-        assert "Total revenue" in call_kwargs["content"]
+        media = call_kwargs.get("media", [])
-        # No media (docx is not an image)
+        assert len(media) == 1
-        assert not call_kwargs.get("media")
+        assert "report.docx" in media[0]
     finally:
         os.chdir(original_cwd)