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

refactor: move document extraction from ContextBuilder to API layer

Browse Source 
ContextBuilder._build_user_content now only handles images (its original
responsibility).  Document text extraction (PDF, DOCX, XLSX, PPTX) is
performed by the new _extract_documents() helper in server.py, called
before process_direct().  This keeps the core context builder free of
format-specific dependencies and makes the API boundary the single place
where uploaded files are pre-processed.

Tests updated to reflect the new responsibility boundary.

Made-with: Cursor
This commit is contained in:
Xubin Ren
2026-04-14 13:00:59 +00:00
parent 2502fc616b
commit 47f5795708
4 changed files with 131 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_api_attachment.py
+60 -8
View File
@@ -10,6 +10,7 @@ import pytest
import pytest_asyncio
from nanobot.api.server import (
_extract_documents,
_FileSizeExceeded,
_parse_json_content,
_save_base64_data_url,
@@ -184,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 saves file to media dir and passes path to process_direct."""
"""Multipart upload of non-image extracts text into content (not media)."""
import os
original_cwd = os.getcwd()
os.chdir(tmp_path)
@@ -202,8 +203,9 @@ 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 call_kwargs["content"] == "analyze this"
assert len(call_kwargs.get("media", [])) == 1
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")
finally:
os.chdir(original_cwd)
@@ -371,13 +373,62 @@ async def test_json_base64_image_upload(aiohttp_client, mock_agent, tmp_path) ->
# ---------------------------------------------------------------------------
# DOCX document extraction tests
# _extract_documents tests (API-layer document extraction)
# ---------------------------------------------------------------------------
def test_extract_documents_separates_images_from_docs(tmp_path) -> None:
"""Images stay in media; document text is appended to content."""
from docx import Document
png = tmp_path / "chart.png"
png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
doc = Document()
doc.add_paragraph("Quarterly revenue is $5M")
docx_path = tmp_path / "report.docx"
doc.save(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
assert "summarize" in text
def test_extract_documents_skips_extraction_errors(tmp_path, monkeypatch) -> None:
"""Document extraction errors should not leak into user text."""
bad_file = tmp_path / "broken.docx"
bad_file.write_text("not a docx", encoding="utf-8")
import nanobot.api.server as _srv
monkeypatch.setattr(
_srv, "extract_text",
lambda _path: "[error: failed to extract DOCX: boom]",
)
text, image_paths = _extract_documents("hello", [str(bad_file)])
assert text == "hello"
assert image_paths == []
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)])
assert text == "describe"
assert len(image_paths) == 1
# ---------------------------------------------------------------------------
# DOCX end-to-end upload test (API layer now 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 should have its text extracted before being sent to AI."""
"""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")
@@ -405,8 +456,9 @@ async def test_docx_upload_extracted_and_sent(aiohttp_client, tmp_path) -> None:
resp = await client.post("/v1/chat/completions", data=data)
assert resp.status == 200
call_kwargs = agent.process_direct.call_args.kwargs
media = call_kwargs.get("media", [])
assert len(media) == 1
assert "report.docx" in media[0]
# 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")
finally:
os.chdir(original_cwd)
tests/test_context_documents.py
+18 -41
View File
@@ -1,4 +1,8 @@
"""Tests for context builder document handling."""
"""Tests for context builder media handling.
The ContextBuilder._build_user_content method should ONLY handle images.
Document text extraction is the responsibility of the API layer.
"""
from __future__ import annotations
@@ -30,52 +34,25 @@ def test_build_user_content_with_image_returns_list(tmp_path: Path) -> None:
assert "text" in types
def test_build_user_content_with_docx_includes_extracted_text(tmp_path: Path) -> None:
"""Document files should have their text extracted and included."""
from docx import Document
doc = Document()
doc.add_paragraph("Quarterly revenue is $5M")
docx_path = tmp_path / "report.docx"
doc.save(docx_path)
def test_build_user_content_ignores_non_image_files(tmp_path: Path) -> None:
"""Non-image files should be silently skipped — extraction is not context builder's job."""
builder = _make_builder(tmp_path)
result = builder._build_user_content("summarize this", [str(docx_path)])
assert isinstance(result, str)
assert "Quarterly revenue" in result
txt = tmp_path / "notes.txt"
txt.write_text("some text", encoding="utf-8")
result = builder._build_user_content("summarize", [str(txt)])
assert result == "summarize"
def test_build_user_content_mixed_image_and_document(tmp_path: Path) -> None:
"""Mix of images and documents: images as base64, docs as text."""
from docx import Document
def test_build_user_content_mixed_image_and_non_image(tmp_path: Path) -> None:
"""Only images should be included; non-image files are skipped."""
builder = _make_builder(tmp_path)
png = tmp_path / "chart.png"
png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
txt = tmp_path / "report.txt"
txt.write_text("report text", encoding="utf-8")
doc = Document()
doc.add_paragraph("Report text here")
docx = tmp_path / "report.docx"
doc.save(docx)
builder = _make_builder(tmp_path)
result = builder._build_user_content("analyze both", [str(png), str(docx)])
result = builder._build_user_content("analyze", [str(png), str(txt)])
assert isinstance(result, list)
assert any(b["type"] == "image_url" for b in result)
text_parts = [b.get("text", "") for b in result if b.get("type") == "text"]
assert any("Report text here" in t for t in text_parts)
def test_build_user_content_skips_document_extraction_errors(tmp_path: Path, monkeypatch) -> None:
"""Document extraction errors should not be embedded into the user prompt."""
docx_path = tmp_path / "broken.docx"
docx_path.write_text("not a real docx", encoding="utf-8")
builder = _make_builder(tmp_path)
monkeypatch.setattr(
"nanobot.utils.document.extract_text",
lambda _path: "[error: failed to extract DOCX: boom]",
)
result = builder._build_user_content("summarize this", [str(docx_path)])
assert result == "summarize this"
assert all("report text" not in t for t in text_parts)