mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-07 21:08:34 +03:00
fix(agent): read document attachments on demand (#5122)
This commit is contained in:
chengyongru
@@ -0,0 +1,208 @@
|
||||
import asyncio
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.loop import AgentLoop, TurnContext, TurnKind
|
||||
from nanobot.agent.tools.filesystem import ReadFileTool
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.config.schema import ChannelsConfig
|
||||
from nanobot.providers.base import LLMResponse
|
||||
from nanobot.utils.document import reference_non_image_attachments
|
||||
|
||||
|
||||
def _make_loop(
|
||||
workspace: Path,
|
||||
channels_config: ChannelsConfig | None = None,
|
||||
) -> AgentLoop:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="ok"))
|
||||
return AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=workspace,
|
||||
model="test-model",
|
||||
channels_config=channels_config,
|
||||
)
|
||||
|
||||
|
||||
def _turn_context(loop: AgentLoop, msg: InboundMessage) -> TurnContext:
|
||||
return TurnContext(
|
||||
msg=msg,
|
||||
session_key=f"{msg.channel}:{msg.chat_id}",
|
||||
turn_id="turn-1",
|
||||
runtime=loop.llm_runtime(),
|
||||
kind=TurnKind.USER,
|
||||
delivery=loop.turn_delivery_factory.create(msg, f"{msg.channel}:{msg.chat_id}"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("extract_document_text", [True, False])
|
||||
async def test_document_attachment_is_referenced_and_read_on_demand(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
extract_document_text: bool,
|
||||
) -> None:
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
media_dir = tmp_path / "media"
|
||||
media_dir.mkdir()
|
||||
csv_path = media_dir / "report.csv"
|
||||
csv_path.write_text("name,value\nnanobot,1", encoding="utf-8")
|
||||
monkeypatch.setattr("nanobot.agent.tools.path_utils.get_media_dir", lambda: media_dir)
|
||||
|
||||
loop = _make_loop(
|
||||
workspace,
|
||||
ChannelsConfig(extract_document_text=extract_document_text),
|
||||
)
|
||||
msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u",
|
||||
chat_id="c",
|
||||
content="import this report",
|
||||
media=[str(csv_path)],
|
||||
)
|
||||
ctx = _turn_context(loop, msg)
|
||||
|
||||
await loop._restore_turn(ctx)
|
||||
|
||||
assert ctx.msg.content == f"import this report\n\n[Attachment: {csv_path}]"
|
||||
assert "name,value" not in ctx.msg.content
|
||||
assert ctx.msg.media == []
|
||||
|
||||
read_tool = ReadFileTool(workspace=workspace, allowed_dir=workspace)
|
||||
result = await read_tool.execute(path=str(csv_path))
|
||||
|
||||
assert "1| name,value" in result
|
||||
assert "2| nanobot,1" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_reference_survives_session_reload(tmp_path: Path) -> None:
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
doc_path = tmp_path / "report.csv"
|
||||
doc_path.write_text("name,value", encoding="utf-8")
|
||||
|
||||
loop = _make_loop(workspace)
|
||||
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("interrupt")) # type: ignore[method-assign]
|
||||
msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u",
|
||||
chat_id="persisted-attachment",
|
||||
content="review this",
|
||||
media=[str(doc_path)],
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="interrupt"):
|
||||
await loop._process_message(msg)
|
||||
|
||||
session_key = "websocket:persisted-attachment"
|
||||
loop.sessions.invalidate(session_key)
|
||||
persisted = loop.sessions.get_or_create(session_key)
|
||||
|
||||
assert [message["role"] for message in persisted.messages] == ["user"]
|
||||
assert persisted.messages[0]["content"] == (
|
||||
f"review this\n\n[Attachment: {doc_path.resolve()}]"
|
||||
)
|
||||
assert "media" not in persisted.messages[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_document_attachment_keeps_body_out_of_prompt(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
doc_path = tmp_path / "followup.txt"
|
||||
doc_path.write_text("Do not inject this file body", encoding="utf-8")
|
||||
captured_messages: list[list[dict]] = []
|
||||
call_count = 0
|
||||
|
||||
async def chat_with_retry(*, messages: list[dict], **kwargs: object) -> LLMResponse:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
captured_messages.append([dict(message) for message in messages])
|
||||
return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage={})
|
||||
|
||||
loop = _make_loop(workspace)
|
||||
loop.provider.chat_with_retry = chat_with_retry
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
|
||||
pending_queue: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
||||
await pending_queue.put(
|
||||
InboundMessage(
|
||||
channel="cli",
|
||||
sender_id="u",
|
||||
chat_id="c",
|
||||
content="check this",
|
||||
media=[str(doc_path)],
|
||||
)
|
||||
)
|
||||
|
||||
final_content, _, _, _, had_injections = await loop._run_agent_loop(
|
||||
[{"role": "user", "content": "hello"}],
|
||||
runtime=loop.llm_runtime(),
|
||||
channel="cli",
|
||||
chat_id="c",
|
||||
pending_queue=pending_queue,
|
||||
)
|
||||
|
||||
assert final_content == "answer-2"
|
||||
assert had_injections is True
|
||||
injected_user_content = [
|
||||
message["content"]
|
||||
for message in captured_messages[-1]
|
||||
if message.get("role") == "user" and isinstance(message.get("content"), str)
|
||||
][-1]
|
||||
assert "check this" in injected_user_content
|
||||
assert f"[Attachment: {doc_path}]" in injected_user_content
|
||||
assert "Do not inject this file body" not in injected_user_content
|
||||
|
||||
|
||||
def test_attachment_references_still_preserve_images(tmp_path: Path) -> None:
|
||||
image_path = tmp_path / "chart.png"
|
||||
image_path.write_bytes(
|
||||
base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+yF9kAAAAASUVORK5CYII="
|
||||
)
|
||||
)
|
||||
doc_path = tmp_path / "report.txt"
|
||||
doc_path.write_text("manual extraction target", encoding="utf-8")
|
||||
|
||||
content, media = reference_non_image_attachments(
|
||||
"review these",
|
||||
[str(image_path), str(doc_path)],
|
||||
)
|
||||
|
||||
assert media == [str(image_path)]
|
||||
assert f"[Attachment: {doc_path}]" in content
|
||||
assert "manual extraction target" not in content
|
||||
|
||||
|
||||
def test_attachment_references_canonicalize_existing_relative_paths(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
image_path = tmp_path / "chart.png"
|
||||
image_path.write_bytes(
|
||||
base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+yF9kAAAAASUVORK5CYII="
|
||||
)
|
||||
)
|
||||
doc_path = tmp_path / "report.csv"
|
||||
doc_path.write_text("name,value", encoding="utf-8")
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
content, media = reference_non_image_attachments(
|
||||
"review these",
|
||||
[image_path.name, doc_path.name],
|
||||
)
|
||||
|
||||
assert media == [str(image_path.resolve())]
|
||||
assert f"[Attachment: {doc_path.resolve()}]" in content
|
||||
@@ -244,38 +244,38 @@ class TestBundledToolContract:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _build_user_content
|
||||
# build_user_content
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBuildUserContent:
|
||||
def test_no_media_returns_string(self, tmp_path):
|
||||
builder = _builder(tmp_path)
|
||||
result = builder._build_user_content("hello", None)
|
||||
result = builder.build_user_content("hello", None)
|
||||
assert result == "hello"
|
||||
|
||||
def test_empty_media_returns_string(self, tmp_path):
|
||||
builder = _builder(tmp_path)
|
||||
result = builder._build_user_content("hello", [])
|
||||
result = builder.build_user_content("hello", [])
|
||||
assert result == "hello"
|
||||
|
||||
def test_nonexistent_media_file_returns_string(self, tmp_path):
|
||||
builder = _builder(tmp_path)
|
||||
result = builder._build_user_content("hello", ["/nonexistent/image.png"])
|
||||
result = builder.build_user_content("hello", ["/nonexistent/image.png"])
|
||||
assert result == "hello"
|
||||
|
||||
def test_non_image_file_returns_string(self, tmp_path):
|
||||
txt = tmp_path / "doc.txt"
|
||||
txt.write_text("not an image", encoding="utf-8")
|
||||
builder = _builder(tmp_path)
|
||||
result = builder._build_user_content("hello", [str(txt)])
|
||||
result = builder.build_user_content("hello", [str(txt)])
|
||||
assert result == "hello"
|
||||
|
||||
def test_valid_image_returns_list(self, tmp_path):
|
||||
png = tmp_path / "test.png"
|
||||
png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 16)
|
||||
builder = _builder(tmp_path)
|
||||
result = builder._build_user_content("hello", [str(png)])
|
||||
result = builder.build_user_content("hello", [str(png)])
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 2
|
||||
assert result[0]["type"] == "image_url"
|
||||
@@ -287,7 +287,7 @@ class TestBuildUserContent:
|
||||
png = tmp_path / "test.png"
|
||||
png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 16)
|
||||
builder = _builder(tmp_path)
|
||||
result = builder._build_user_content("hello", [str(png)])
|
||||
result = builder.build_user_content("hello", [str(png)])
|
||||
assert "_meta" in result[0]
|
||||
assert "path" in result[0]["_meta"]
|
||||
|
||||
|
||||
@@ -1,176 +0,0 @@
|
||||
import asyncio
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.loop import AgentLoop, TurnContext, TurnKind
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.config.schema import ChannelsConfig
|
||||
from nanobot.providers.base import LLMResponse
|
||||
from nanobot.utils.document import reference_non_image_attachments
|
||||
|
||||
|
||||
def _make_loop(tmp_path: Path, channels_config: ChannelsConfig | None = None) -> AgentLoop:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="ok"))
|
||||
return AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
channels_config=channels_config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_turn_extracts_documents_by_default(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop = _make_loop(tmp_path)
|
||||
doc_path = tmp_path / "report.txt"
|
||||
doc_path.write_text("Quarterly revenue is $5M", encoding="utf-8")
|
||||
calls: list[tuple[str, list[str]]] = []
|
||||
|
||||
def fake_extract_documents(content: str, media: list[str]) -> tuple[str, list[str]]:
|
||||
calls.append((content, media))
|
||||
return f"{content}\n\n[File: report.txt]\nQuarterly revenue is $5M", []
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.loop.extract_documents", fake_extract_documents)
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="cli",
|
||||
sender_id="u",
|
||||
chat_id="c",
|
||||
content="summarize",
|
||||
media=[str(doc_path)],
|
||||
)
|
||||
ctx = TurnContext(
|
||||
msg=msg,
|
||||
session_key="cli:c",
|
||||
turn_id="turn-1",
|
||||
runtime=loop.llm_runtime(),
|
||||
kind=TurnKind.USER,
|
||||
delivery=loop.turn_delivery_factory.create(msg, "cli:c"),
|
||||
)
|
||||
|
||||
await loop._restore_turn(ctx)
|
||||
|
||||
assert calls == [("summarize", [str(doc_path)])]
|
||||
assert "Quarterly revenue" in ctx.msg.content
|
||||
assert ctx.msg.media == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_turn_references_documents_when_extraction_disabled(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
loop = _make_loop(tmp_path, ChannelsConfig(extract_document_text=False))
|
||||
doc_path = tmp_path / "report.txt"
|
||||
doc_path.write_text("Quarterly revenue is $5M", encoding="utf-8")
|
||||
|
||||
def fail_extract_documents(content: str, media: list[str]) -> tuple[str, list[str]]:
|
||||
raise AssertionError("document extraction should be disabled")
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.loop.extract_documents", fail_extract_documents)
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="cli",
|
||||
sender_id="u",
|
||||
chat_id="c",
|
||||
content="summarize",
|
||||
media=[str(doc_path)],
|
||||
)
|
||||
ctx = TurnContext(
|
||||
msg=msg,
|
||||
session_key="cli:c",
|
||||
turn_id="turn-1",
|
||||
runtime=loop.llm_runtime(),
|
||||
kind=TurnKind.USER,
|
||||
delivery=loop.turn_delivery_factory.create(msg, "cli:c"),
|
||||
)
|
||||
|
||||
await loop._restore_turn(ctx)
|
||||
|
||||
assert "Quarterly revenue" not in ctx.msg.content
|
||||
assert f"[Attachment: {doc_path}]" in ctx.msg.content
|
||||
assert ctx.msg.media == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_followup_references_documents_when_extraction_disabled(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
doc_path = tmp_path / "followup.txt"
|
||||
doc_path.write_text("Do not inject this file body", encoding="utf-8")
|
||||
captured_messages: list[list[dict]] = []
|
||||
call_count = {"n": 0}
|
||||
|
||||
async def chat_with_retry(*, messages: list[dict], **kwargs: object) -> LLMResponse:
|
||||
call_count["n"] += 1
|
||||
captured_messages.append([dict(message) for message in messages])
|
||||
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage={})
|
||||
|
||||
loop = _make_loop(tmp_path, ChannelsConfig(extract_document_text=False))
|
||||
loop.provider.chat_with_retry = chat_with_retry
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
|
||||
def fail_extract_documents(content: str, media: list[str]) -> tuple[str, list[str]]:
|
||||
raise AssertionError("document extraction should be disabled")
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.loop.extract_documents", fail_extract_documents)
|
||||
|
||||
pending_queue: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
||||
await pending_queue.put(
|
||||
InboundMessage(
|
||||
channel="cli",
|
||||
sender_id="u",
|
||||
chat_id="c",
|
||||
content="check this",
|
||||
media=[str(doc_path)],
|
||||
)
|
||||
)
|
||||
|
||||
final_content, _, _, _, had_injections = await loop._run_agent_loop(
|
||||
[{"role": "user", "content": "hello"}],
|
||||
runtime=loop.llm_runtime(),
|
||||
channel="cli",
|
||||
chat_id="c",
|
||||
pending_queue=pending_queue,
|
||||
)
|
||||
|
||||
assert final_content == "answer-2"
|
||||
assert had_injections is True
|
||||
injected_user_content = [
|
||||
message["content"]
|
||||
for message in captured_messages[-1]
|
||||
if message.get("role") == "user" and isinstance(message.get("content"), str)
|
||||
][-1]
|
||||
assert "check this" in injected_user_content
|
||||
assert f"[Attachment: {doc_path}]" in injected_user_content
|
||||
assert "Do not inject this file body" not in injected_user_content
|
||||
|
||||
|
||||
def test_document_extraction_disabled_still_preserves_images(tmp_path: Path) -> None:
|
||||
image_path = tmp_path / "chart.png"
|
||||
image_path.write_bytes(
|
||||
base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+yF9kAAAAASUVORK5CYII="
|
||||
)
|
||||
)
|
||||
doc_path = tmp_path / "report.txt"
|
||||
doc_path.write_text("manual extraction target", encoding="utf-8")
|
||||
|
||||
content, media = reference_non_image_attachments(
|
||||
"review these",
|
||||
[str(image_path), str(doc_path)],
|
||||
)
|
||||
|
||||
assert media == [str(image_path)]
|
||||
assert f"[Attachment: {doc_path}]" in content
|
||||
@@ -713,10 +713,9 @@ def test_unified_session_route_ignores_non_user_destinations(
|
||||
assert session.metadata[LAST_CHANNEL_METADATA_KEY] == "telegram:existing"
|
||||
|
||||
|
||||
# 1x1 PNG used by the media-persistence tests. ``extract_documents`` runs
|
||||
# at the top of ``_process_message`` and filters ``msg.media`` down to
|
||||
# paths that magic-byte-sniff as images, so the test fixture needs real
|
||||
# bytes on disk (not just placeholder paths).
|
||||
# 1x1 PNG used by the media-persistence tests. Attachment preparation filters
|
||||
# ``msg.media`` down to paths that magic-byte-sniff as images, so the test
|
||||
# fixture needs real bytes on disk (not just placeholder paths).
|
||||
_PNG_1X1 = (
|
||||
b"\x89PNG\r\n\x1a\n"
|
||||
b"\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
|
||||
|
||||
Reference in New Issue
Block a user