Add Ollama image generation support

This commit is contained in:
Haisam Abbas
2026-05-21 12:06:08 +05:00
parent eae51333ad
commit 84603f4cf2
5 changed files with 259 additions and 4 deletions
+49
View File
@@ -13,6 +13,7 @@ from nanobot.providers.image_generation import (
GeneratedImageResponse,
ImageGenerationError,
MiniMaxImageGenerationClient,
OllamaImageGenerationClient,
OpenRouterImageGenerationClient,
StepFunImageGenerationClient,
)
@@ -133,6 +134,54 @@ async def test_openrouter_image_generation_requires_api_key() -> None:
await client.generate(prompt="draw", model="model")
@pytest.mark.asyncio
async def test_ollama_image_generation_payload_and_response() -> None:
raw_b64 = PNG_DATA_URL.removeprefix("data:image/png;base64,")
fake = FakeClient(FakeResponse({"image": raw_b64}))
client = OllamaImageGenerationClient(
api_key="ollama-test",
api_base="http://localhost:11434/v1/",
extra_headers={"X-Test": "1"},
extra_body={"seed": 123},
client=fake, # type: ignore[arg-type]
)
response = await client.generate(
prompt="a sunset",
model="x/z-image-turbo",
aspect_ratio="16:9",
image_size="1K",
)
assert response.images == [PNG_DATA_URL]
assert response.content == ""
call = fake.calls[0]
assert call["url"] == "http://localhost:11434/api/generate"
assert call["headers"]["Authorization"] == "Bearer ollama-test"
assert call["headers"]["X-Test"] == "1"
body = call["json"]
assert body["model"] == "x/z-image-turbo"
assert body["prompt"] == "a sunset"
assert body["width"] == 1024
assert body["height"] == 576
assert body["steps"] == 0
assert body["stream"] is False
assert body["seed"] == 123
@pytest.mark.asyncio
async def test_ollama_image_generation_rejects_reference_images() -> None:
client = OllamaImageGenerationClient(api_key=None)
with pytest.raises(ImageGenerationError, match="reference images"):
await client.generate(
prompt="edit this",
model="x/z-image-turbo",
reference_images=["ref.png"],
)
@pytest.mark.asyncio
async def test_aihubmix_image_generation_payload_and_response() -> None:
raw_b64 = PNG_DATA_URL.removeprefix("data:image/png;base64,")
+33
View File
@@ -138,6 +138,39 @@ async def test_generate_image_tool_reports_missing_aihubmix_key(tmp_path: Path)
assert result.startswith("Error: AIHubMix API key is not configured")
@pytest.mark.asyncio
async def test_generate_image_tool_allows_ollama_without_api_key(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
set_config_path(tmp_path / "config.json")
FakeImageClient.instances = []
monkeypatch.setattr(
"nanobot.agent.tools.image_generation.get_image_gen_provider",
lambda name: FakeImageClient if name == "ollama" else None,
)
tool = ImageGenerationTool(
workspace=tmp_path,
config=ImageGenerationToolConfig(
enabled=True,
provider="ollama",
model="x/z-image-turbo",
),
provider_configs={"ollama": ProviderConfig(api_base="http://localhost:11434/v1")},
)
result = await tool.execute(prompt="draw a cat")
payload = json.loads(result)
assert len(payload["artifacts"]) == 1
fake = FakeImageClient.instances[0]
assert fake.kwargs["api_key"] is None
assert fake.kwargs["api_base"] == "http://localhost:11434/v1"
assert fake.calls[0]["aspect_ratio"] == "1:1"
assert fake.calls[0]["image_size"] == "1K"
@pytest.mark.asyncio
async def test_generate_image_tool_rejects_reference_outside_workspace(tmp_path: Path) -> None:
set_config_path(tmp_path / "config.json")