feat: add image generation tool and WebUI mode
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xubin Ren
committed by
Xubin Ren
Xubin Ren
co-authored by
Cursor
Cursor
parent
3a2f47d720
commit
e936ed48bd
@@ -0,0 +1,89 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.config.loader import set_config_path
|
||||
from nanobot.config.schema import ImageGenerationToolConfig, ProviderConfig, ToolsConfig
|
||||
from nanobot.providers.base import LLMResponse, ToolCallRequest
|
||||
from nanobot.providers.image_generation import GeneratedImageResponse
|
||||
|
||||
PNG_DATA_URL = (
|
||||
"data:image/png;base64,"
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII="
|
||||
)
|
||||
|
||||
|
||||
class FakeImageClient:
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
async def generate(self, **kwargs: Any) -> GeneratedImageResponse:
|
||||
return GeneratedImageResponse(images=[PNG_DATA_URL], content="", raw={})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_image_media_is_attached_to_final_assistant_message(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
set_config_path(tmp_path / "config.json")
|
||||
monkeypatch.setattr(
|
||||
"nanobot.agent.tools.image_generation.OpenRouterImageGenerationClient",
|
||||
FakeImageClient,
|
||||
)
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.generation.max_tokens = 4096
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
side_effect=[
|
||||
LLMResponse(
|
||||
content="",
|
||||
finish_reason="tool_calls",
|
||||
tool_calls=[
|
||||
ToolCallRequest(
|
||||
id="call_img",
|
||||
name="generate_image",
|
||||
arguments={"prompt": "draw a tiny icon"},
|
||||
)
|
||||
],
|
||||
),
|
||||
LLMResponse(content="Done", finish_reason="stop"),
|
||||
]
|
||||
)
|
||||
provider.chat_stream_with_retry = AsyncMock()
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
tools_config=ToolsConfig(
|
||||
image_generation=ImageGenerationToolConfig(enabled=True),
|
||||
),
|
||||
image_generation_provider_config=ProviderConfig(api_key="sk-or-test"),
|
||||
)
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
result = await loop._process_message(
|
||||
InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="chat-image",
|
||||
content="draw an icon",
|
||||
)
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.content == "Done"
|
||||
assert len(result.media) == 1
|
||||
assert Path(result.media[0]).is_file()
|
||||
|
||||
session = loop.sessions.get_or_create("websocket:chat-image")
|
||||
assert session.messages[-1]["role"] == "assistant"
|
||||
assert session.messages[-1]["media"] == result.media
|
||||
@@ -0,0 +1,204 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from nanobot.providers.image_generation import (
|
||||
AIHubMixImageGenerationClient,
|
||||
GeneratedImageResponse,
|
||||
ImageGenerationError,
|
||||
OpenRouterImageGenerationClient,
|
||||
)
|
||||
|
||||
PNG_BYTES = (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
|
||||
b"\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02"
|
||||
b"\x00\x00\x00\x0bIDATx\xdacd\xfc\xff\x1f\x00\x03\x03"
|
||||
b"\x02\x00\xef\xbf\xa7\xdb\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
PNG_DATA_URL = (
|
||||
"data:image/png;base64,"
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII="
|
||||
)
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(
|
||||
self,
|
||||
payload: dict[str, Any],
|
||||
status_code: int = 200,
|
||||
content: bytes = b"",
|
||||
) -> None:
|
||||
self._payload = payload
|
||||
self.status_code = status_code
|
||||
self.text = str(payload)
|
||||
self.content = content
|
||||
self.request = httpx.Request("POST", "https://openrouter.ai/api/v1/chat/completions")
|
||||
|
||||
def json(self) -> dict[str, Any]:
|
||||
return self._payload
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
if self.status_code >= 400:
|
||||
response = httpx.Response(self.status_code, request=self.request, text=self.text)
|
||||
raise httpx.HTTPStatusError("failed", request=self.request, response=response)
|
||||
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, response: FakeResponse) -> None:
|
||||
self.response = response
|
||||
self.get_response = response
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.get_calls: list[dict[str, Any]] = []
|
||||
|
||||
async def post(self, url: str, **kwargs: Any) -> FakeResponse:
|
||||
self.calls.append({"url": url, **kwargs})
|
||||
return self.response
|
||||
|
||||
async def get(self, url: str, **kwargs: Any) -> FakeResponse:
|
||||
self.get_calls.append({"url": url, **kwargs})
|
||||
return self.get_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_image_generation_payload_and_response(tmp_path: Path) -> None:
|
||||
ref = tmp_path / "ref.png"
|
||||
ref.write_bytes(PNG_BYTES)
|
||||
fake = FakeClient(
|
||||
FakeResponse(
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": "done",
|
||||
"images": [{"image_url": {"url": PNG_DATA_URL}}],
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
client = OpenRouterImageGenerationClient(
|
||||
api_key="sk-or-test",
|
||||
api_base="https://openrouter.ai/api/v1/",
|
||||
extra_headers={"X-Test": "1"},
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(
|
||||
prompt="make this blue",
|
||||
model="openai/gpt-5.4-image-2",
|
||||
reference_images=[str(ref)],
|
||||
aspect_ratio="16:9",
|
||||
image_size="2K",
|
||||
)
|
||||
|
||||
assert isinstance(response, GeneratedImageResponse)
|
||||
assert response.images == [PNG_DATA_URL]
|
||||
assert response.content == "done"
|
||||
|
||||
call = fake.calls[0]
|
||||
assert call["url"] == "https://openrouter.ai/api/v1/chat/completions"
|
||||
assert call["headers"]["Authorization"] == "Bearer sk-or-test"
|
||||
assert call["headers"]["X-Test"] == "1"
|
||||
body = call["json"]
|
||||
assert body["modalities"] == ["image", "text"]
|
||||
assert body["image_config"] == {"aspect_ratio": "16:9", "image_size": "2K"}
|
||||
assert body["messages"][0]["content"][0] == {"type": "text", "text": "make this blue"}
|
||||
assert body["messages"][0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_image_generation_requires_images() -> None:
|
||||
fake = FakeClient(FakeResponse({"choices": [{"message": {"content": "text only"}}]}))
|
||||
client = OpenRouterImageGenerationClient(api_key="sk-or-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="returned no images"):
|
||||
await client.generate(prompt="draw", model="model")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_image_generation_requires_api_key() -> None:
|
||||
client = OpenRouterImageGenerationClient(api_key=None)
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="API key"):
|
||||
await client.generate(prompt="draw", model="model")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aihubmix_image_generation_payload_and_response() -> None:
|
||||
raw_b64 = PNG_DATA_URL.removeprefix("data:image/png;base64,")
|
||||
fake = FakeClient(FakeResponse({"output": {"b64_json": [{"bytesBase64": raw_b64}]}}))
|
||||
client = AIHubMixImageGenerationClient(
|
||||
api_key="sk-ahm-test",
|
||||
api_base="https://aihubmix.com/v1/",
|
||||
extra_headers={"APP-Code": "nanobot"},
|
||||
extra_body={"quality": "low"},
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(
|
||||
prompt="draw a logo",
|
||||
model="gpt-image-2-free",
|
||||
aspect_ratio="16:9",
|
||||
image_size="1K",
|
||||
)
|
||||
|
||||
assert response.images == [PNG_DATA_URL]
|
||||
call = fake.calls[0]
|
||||
assert call["url"] == "https://aihubmix.com/v1/models/openai/gpt-image-2-free/predictions"
|
||||
assert call["headers"]["Authorization"] == "Bearer sk-ahm-test"
|
||||
assert call["headers"]["APP-Code"] == "nanobot"
|
||||
assert call["json"] == {
|
||||
"input": {
|
||||
"prompt": "draw a logo",
|
||||
"n": 1,
|
||||
"size": "1536x1024",
|
||||
"quality": "low",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aihubmix_image_edit_payload_uses_reference_images(tmp_path: Path) -> None:
|
||||
raw_b64 = PNG_DATA_URL.removeprefix("data:image/png;base64,")
|
||||
fake = FakeClient(FakeResponse({"output": [{"b64_json": raw_b64}]}))
|
||||
ref = tmp_path / "ref.png"
|
||||
ref.write_bytes(PNG_BYTES)
|
||||
client = AIHubMixImageGenerationClient(
|
||||
api_key="sk-ahm-test",
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(
|
||||
prompt="edit this",
|
||||
model="gpt-image-2-free",
|
||||
reference_images=[str(ref)],
|
||||
aspect_ratio="1:1",
|
||||
)
|
||||
|
||||
assert response.images == [PNG_DATA_URL]
|
||||
call = fake.calls[0]
|
||||
assert call["url"] == "https://aihubmix.com/v1/models/openai/gpt-image-2-free/predictions"
|
||||
assert call["json"]["input"]["prompt"] == "edit this"
|
||||
assert call["json"]["input"]["n"] == 1
|
||||
assert call["json"]["input"]["size"] == "1024x1024"
|
||||
assert call["json"]["input"]["image"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aihubmix_image_generation_downloads_url_response() -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
|
||||
fake.get_response = FakeResponse({}, content=PNG_BYTES)
|
||||
client = AIHubMixImageGenerationClient(
|
||||
api_key="sk-ahm-test",
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(prompt="draw", model="gpt-image-2-free")
|
||||
|
||||
assert response.images[0].startswith("data:image/png;base64,")
|
||||
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
|
||||
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.image_generation import ImageGenerationTool
|
||||
from nanobot.config.loader import set_config_path
|
||||
from nanobot.config.schema import ImageGenerationToolConfig, ProviderConfig
|
||||
from nanobot.providers.image_generation import GeneratedImageResponse
|
||||
|
||||
PNG_BYTES = (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
|
||||
b"\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02"
|
||||
b"\x00\x00\x00\x0bIDATx\xdacd\xfc\xff\x1f\x00\x03\x03"
|
||||
b"\x02\x00\xef\xbf\xa7\xdb\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
PNG_DATA_URL = (
|
||||
"data:image/png;base64,"
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII="
|
||||
)
|
||||
|
||||
|
||||
class FakeImageClient:
|
||||
instances: list["FakeImageClient"] = []
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self.kwargs = kwargs
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.instances.append(self)
|
||||
|
||||
async def generate(self, **kwargs: Any) -> GeneratedImageResponse:
|
||||
self.calls.append(kwargs)
|
||||
return GeneratedImageResponse(images=[PNG_DATA_URL], content="", raw={})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_tool_stores_artifact_and_source_images(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
set_config_path(tmp_path / "config.json")
|
||||
FakeImageClient.instances = []
|
||||
monkeypatch.setattr(
|
||||
"nanobot.agent.tools.image_generation.OpenRouterImageGenerationClient",
|
||||
FakeImageClient,
|
||||
)
|
||||
ref = tmp_path / "ref.png"
|
||||
ref.write_bytes(PNG_BYTES)
|
||||
tool = ImageGenerationTool(
|
||||
workspace=tmp_path,
|
||||
config=ImageGenerationToolConfig(enabled=True, max_images_per_turn=2),
|
||||
provider_config=ProviderConfig(api_key="sk-or-test"),
|
||||
)
|
||||
|
||||
result = await tool.execute(
|
||||
prompt="make this blue",
|
||||
reference_images=["ref.png"],
|
||||
aspect_ratio="16:9",
|
||||
image_size="2K",
|
||||
count=2,
|
||||
)
|
||||
|
||||
payload = json.loads(result)
|
||||
artifacts = payload["artifacts"]
|
||||
assert len(artifacts) == 2
|
||||
assert Path(artifacts[0]["path"]).is_file()
|
||||
assert artifacts[0]["source_images"] == [str(ref.resolve())]
|
||||
assert artifacts[0]["model"] == "openai/gpt-5.4-image-2"
|
||||
|
||||
fake = FakeImageClient.instances[0]
|
||||
assert fake.kwargs["api_key"] == "sk-or-test"
|
||||
assert len(fake.calls) == 2
|
||||
assert fake.calls[0]["aspect_ratio"] == "16:9"
|
||||
assert fake.calls[0]["image_size"] == "2K"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_tool_reports_missing_key(tmp_path: Path) -> None:
|
||||
tool = ImageGenerationTool(
|
||||
workspace=tmp_path,
|
||||
config=ImageGenerationToolConfig(enabled=True),
|
||||
provider_config=ProviderConfig(),
|
||||
)
|
||||
|
||||
result = await tool.execute(prompt="draw")
|
||||
|
||||
assert result.startswith("Error: OpenRouter API key is not configured")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_tool_selects_aihubmix_provider(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
set_config_path(tmp_path / "config.json")
|
||||
FakeImageClient.instances = []
|
||||
monkeypatch.setattr(
|
||||
"nanobot.agent.tools.image_generation.AIHubMixImageGenerationClient",
|
||||
FakeImageClient,
|
||||
)
|
||||
tool = ImageGenerationTool(
|
||||
workspace=tmp_path,
|
||||
config=ImageGenerationToolConfig(
|
||||
enabled=True,
|
||||
provider="aihubmix",
|
||||
model="gpt-image-2-free",
|
||||
),
|
||||
provider_configs={
|
||||
"openrouter": ProviderConfig(api_key="sk-or-test"),
|
||||
"aihubmix": ProviderConfig(api_key="sk-ahm-test", extra_body={"quality": "low"}),
|
||||
},
|
||||
)
|
||||
|
||||
result = await tool.execute(prompt="draw a poster", aspect_ratio="3:4")
|
||||
|
||||
payload = json.loads(result)
|
||||
assert len(payload["artifacts"]) == 1
|
||||
fake = FakeImageClient.instances[0]
|
||||
assert fake.kwargs["api_key"] == "sk-ahm-test"
|
||||
assert fake.kwargs["extra_body"] == {"quality": "low"}
|
||||
assert fake.calls[0]["model"] == "gpt-image-2-free"
|
||||
assert fake.calls[0]["aspect_ratio"] == "3:4"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_tool_reports_missing_aihubmix_key(tmp_path: Path) -> None:
|
||||
tool = ImageGenerationTool(
|
||||
workspace=tmp_path,
|
||||
config=ImageGenerationToolConfig(enabled=True, provider="aihubmix"),
|
||||
provider_configs={"aihubmix": ProviderConfig()},
|
||||
)
|
||||
|
||||
result = await tool.execute(prompt="draw")
|
||||
|
||||
assert result.startswith("Error: AIHubMix API key is not configured")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_image_tool_rejects_reference_outside_workspace(tmp_path: Path) -> None:
|
||||
set_config_path(tmp_path / "config.json")
|
||||
outside = tmp_path.parent / "outside.png"
|
||||
outside.write_bytes(PNG_BYTES)
|
||||
tool = ImageGenerationTool(
|
||||
workspace=tmp_path,
|
||||
config=ImageGenerationToolConfig(enabled=True),
|
||||
provider_config=ProviderConfig(api_key="sk-or-test"),
|
||||
)
|
||||
|
||||
result = await tool.execute(prompt="edit", reference_images=[str(outside)])
|
||||
|
||||
assert "reference_images must be inside the workspace" in result
|
||||
@@ -55,6 +55,25 @@ async def test_message_tool_marks_channel_delivery_only_when_enabled() -> None:
|
||||
assert sent[1].metadata == {"_record_channel_delivery": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_tool_records_media_deliveries() -> None:
|
||||
sent: list[OutboundMessage] = []
|
||||
|
||||
async def _send(msg: OutboundMessage) -> None:
|
||||
sent.append(msg)
|
||||
|
||||
tool = MessageTool(send_callback=_send)
|
||||
|
||||
await tool.execute(
|
||||
content="image",
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
media=["/tmp/generated.png"],
|
||||
)
|
||||
|
||||
assert sent[0].metadata == {"_record_channel_delivery": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_tool_inherits_metadata_for_same_target() -> None:
|
||||
sent: list[OutboundMessage] = []
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
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"},
|
||||
]
|
||||
)
|
||||
|
||||
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"]
|
||||
@@ -0,0 +1,25 @@
|
||||
from nanobot.utils.image_generation_intent import image_generation_prompt
|
||||
|
||||
|
||||
def test_image_generation_prompt_ignores_plain_messages() -> None:
|
||||
assert image_generation_prompt("hello", {}) == "hello"
|
||||
|
||||
|
||||
def test_image_generation_prompt_uses_auto_aspect_instruction() -> None:
|
||||
prompt = image_generation_prompt(
|
||||
"Draw a poster",
|
||||
{"image_generation": {"enabled": True, "aspect_ratio": None}},
|
||||
)
|
||||
|
||||
assert "Draw a poster" in prompt
|
||||
assert "Use the generate_image tool" in prompt
|
||||
assert "Choose the most suitable aspect_ratio yourself" in prompt
|
||||
|
||||
|
||||
def test_image_generation_prompt_uses_selected_aspect_ratio() -> None:
|
||||
prompt = image_generation_prompt(
|
||||
"Draw a banner",
|
||||
{"image_generation": {"enabled": True, "aspect_ratio": "16:9"}},
|
||||
)
|
||||
|
||||
assert "aspect_ratio='16:9'" in prompt
|
||||
Reference in New Issue
Block a user