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:
+1
-40
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user