mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-05-19 16:12:30 +00:00
86858cfcb8
The runtime media-attachment mechanism was broken for streaming channels (e.g. WebSocket): the _streamed flag caused _send_once to skip the final OutboundMessage that carried generated media, so images were never delivered. Rather than adding complex coordination between streaming and media delivery, delegate image delivery to the LLM: after generate_image returns artifact paths, the next_step prompt now instructs the LLM to call the message tool with the paths in the media parameter. This works uniformly across all channels, streaming or not. Remove generated_media from TurnContext, _assemble_outbound, and _state_save. Update prompts in identity.md, SKILL.md, message tool description, and artifacts.py to reflect the new flow.
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from nanobot.config.loader import set_config_path
|
|
from nanobot.utils.artifacts import (
|
|
ArtifactError,
|
|
decode_image_data_url,
|
|
generated_image_paths_from_messages,
|
|
generated_image_tool_result,
|
|
store_generated_image_artifact,
|
|
)
|
|
|
|
PNG_DATA_URL = (
|
|
"data:image/png;base64,"
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII="
|
|
)
|
|
|
|
|
|
def test_decode_image_data_url_validates_image_payload() -> None:
|
|
raw, mime = decode_image_data_url(PNG_DATA_URL)
|
|
|
|
assert raw.startswith(b"\x89PNG")
|
|
assert mime == "image/png"
|
|
|
|
with pytest.raises(ArtifactError):
|
|
decode_image_data_url("data:image/png;base64,not-base64")
|
|
|
|
|
|
def test_store_generated_image_artifact_writes_image_and_sidecar(tmp_path: Path) -> None:
|
|
set_config_path(tmp_path / "config.json")
|
|
created_at = datetime(2026, 5, 8, 12, 0, tzinfo=timezone.utc)
|
|
|
|
artifact = store_generated_image_artifact(
|
|
PNG_DATA_URL,
|
|
prompt="draw a tiny pixel",
|
|
model="openai/gpt-5.4-image-2",
|
|
source_images=["/tmp/ref.png"],
|
|
save_dir="generated",
|
|
created_at=created_at,
|
|
)
|
|
|
|
image_path = Path(artifact["path"])
|
|
assert image_path.is_file()
|
|
assert image_path.parent == tmp_path / "media" / "generated" / "2026-05-08"
|
|
assert artifact["id"].startswith("img_")
|
|
assert artifact["mime"] == "image/png"
|
|
|
|
sidecar = image_path.with_suffix(".json")
|
|
metadata = json.loads(sidecar.read_text(encoding="utf-8"))
|
|
assert metadata["path"] == str(image_path)
|
|
assert metadata["source_images"] == ["/tmp/ref.png"]
|
|
|
|
|
|
def test_store_generated_image_artifact_rejects_unsafe_save_dir(tmp_path: Path) -> None:
|
|
set_config_path(tmp_path / "config.json")
|
|
|
|
with pytest.raises(ArtifactError):
|
|
store_generated_image_artifact(
|
|
PNG_DATA_URL,
|
|
prompt="x",
|
|
model="m",
|
|
save_dir="../outside",
|
|
)
|
|
|
|
|
|
def test_generated_image_paths_from_tool_results() -> None:
|
|
result = generated_image_tool_result(
|
|
[
|
|
{"id": "img_1", "path": "/tmp/one.png"},
|
|
{"id": "img_2", "path": "/tmp/two.png"},
|
|
]
|
|
)
|
|
payload = json.loads(result)
|
|
|
|
assert generated_image_paths_from_messages(
|
|
[
|
|
{"role": "tool", "name": "generate_image", "content": result},
|
|
{"role": "tool", "name": "other", "content": result},
|
|
]
|
|
) == ["/tmp/one.png", "/tmp/two.png"]
|
|
assert "Call the message tool" in payload["next_step"]
|
|
assert "media parameter" in payload["next_step"]
|