diff --git a/nanobot/agent/context.py b/nanobot/agent/context.py index 0996def5..cab7b057 100644 --- a/nanobot/agent/context.py +++ b/nanobot/agent/context.py @@ -147,56 +147,30 @@ class ContextBuilder: messages.append({"role": current_role, "content": merged}) return messages - def _build_user_content( - self, text: str, media: list[str] | None - ) -> str | list[dict[str, Any]]: - """Build user message content with optional media. - - Images are converted to base64 vision blocks. - Documents (PDF, Word, Excel, PPT) have their text extracted and appended. - """ + def _build_user_content(self, text: str, media: list[str] | None) -> str | list[dict[str, Any]]: + """Build user message content with optional base64-encoded images.""" if not media: return text - images: list[dict[str, Any]] = [] - doc_texts: list[str] = [] - + images = [] for path in media: p = Path(path) if not p.is_file(): continue raw = p.read_bytes() mime = detect_image_mime(raw) or mimetypes.guess_type(path)[0] + if not mime or not mime.startswith("image/"): + continue + b64 = base64.b64encode(raw).decode() + images.append({ + "type": "image_url", + "image_url": {"url": f"data:{mime};base64,{b64}"}, + "_meta": {"path": str(p)}, + }) - if mime and mime.startswith("image/"): - b64 = base64.b64encode(raw).decode() - images.append({ - "type": "image_url", - "image_url": {"url": f"data:{mime};base64,{b64}"}, - "_meta": {"path": str(p)}, - }) - else: - # Try document text extraction - from nanobot.utils.document import extract_text - extracted = extract_text(p) - if extracted and not extracted.startswith("[error:"): - doc_texts.append(f"[File: {p.name}]\n{extracted}") - - # Build final content - parts: list[dict[str, Any]] = [] - parts.extend(images) - - combined_text = text - if doc_texts: - combined_text = text + "\n\n" + "\n\n".join(doc_texts) - - if images: - parts.append({"type": "text", "text": combined_text}) - return parts - elif doc_texts: - return combined_text - else: + if not images: return text + return images + [{"type": "text", "text": text}] def add_tool_result( self, messages: list[dict[str, Any]], diff --git a/nanobot/api/server.py b/nanobot/api/server.py index 934879a3..fcba8b55 100644 --- a/nanobot/api/server.py +++ b/nanobot/api/server.py @@ -19,7 +19,8 @@ from aiohttp import web from loguru import logger from nanobot.config.paths import get_media_dir -from nanobot.utils.helpers import safe_filename +from nanobot.utils.document import extract_text +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 @@ -161,6 +162,40 @@ 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 # --------------------------------------------------------------------------- @@ -197,6 +232,10 @@ 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/tests/test_api_attachment.py b/tests/test_api_attachment.py index 082494b7..ea8eed9d 100644 --- a/tests/test_api_attachment.py +++ b/tests/test_api_attachment.py @@ -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) diff --git a/tests/test_context_documents.py b/tests/test_context_documents.py index 9f503906..28a4f6d2 100644 --- a/tests/test_context_documents.py +++ b/tests/test_context_documents.py @@ -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)