Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

refactor: centralize document extraction in AgentLoop._process_message

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+7
View File
@@ -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
+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())
nanobot/utils/document.py
+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
tests/test_api_attachment.py
+31 -25
View File
@@ -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)