Add Ollama image generation support
This commit is contained in:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user