refactor(agent): remove legacy image mode prompt injection
This commit is contained in:
@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any
|
||||
from pydantic import Field
|
||||
|
||||
from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters
|
||||
from nanobot.agent.tools.context import RequestContext
|
||||
from nanobot.agent.tools.schema import (
|
||||
ArraySchema,
|
||||
IntegerSchema,
|
||||
@@ -22,7 +21,6 @@ from nanobot.providers.image_generation import (
|
||||
ImageGenerationProvider,
|
||||
get_image_gen_provider,
|
||||
)
|
||||
from nanobot.runtime_context import RuntimeContextBlock
|
||||
from nanobot.security.workspace_access import current_tool_workspace
|
||||
from nanobot.security.workspace_policy import WorkspaceBoundaryError, resolve_allowed_path
|
||||
from nanobot.utils.artifacts import (
|
||||
@@ -31,7 +29,6 @@ from nanobot.utils.artifacts import (
|
||||
store_generated_image_artifact,
|
||||
)
|
||||
from nanobot.utils.helpers import detect_image_mime
|
||||
from nanobot.utils.image_generation_intent import image_generation_runtime_context
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.config.schema import ProviderConfig
|
||||
@@ -119,18 +116,6 @@ class ImageGenerationTool(Tool):
|
||||
"or user image paths as reference_images."
|
||||
)
|
||||
|
||||
def runtime_context_provider(self):
|
||||
return self._provide_runtime_context
|
||||
|
||||
async def _provide_runtime_context(
|
||||
self,
|
||||
request: RequestContext,
|
||||
) -> RuntimeContextBlock | None:
|
||||
content = image_generation_runtime_context(request.metadata)
|
||||
if not content:
|
||||
return None
|
||||
return RuntimeContextBlock(source="image_generation", content=content)
|
||||
|
||||
def _provider_config(self) -> ProviderConfig | None:
|
||||
return self.provider_configs.get(self.config.provider)
|
||||
|
||||
|
||||
@@ -782,13 +782,6 @@ class WebSocketChannel(BaseChannel):
|
||||
metadata["mcp_presets"] = mcp_presets
|
||||
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
|
||||
self._workspaces.persist_scope(cid, scope)
|
||||
image_generation = envelope.get("image_generation")
|
||||
if isinstance(image_generation, dict) and image_generation.get("enabled") is True:
|
||||
aspect_ratio = image_generation.get("aspect_ratio")
|
||||
metadata["image_generation"] = {
|
||||
"enabled": True,
|
||||
"aspect_ratio": aspect_ratio if isinstance(aspect_ratio, str) else None,
|
||||
}
|
||||
if metadata.get("webui") is True and self.is_allowed(client_id):
|
||||
self._transcripts.append_user_message(
|
||||
cid,
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
"""Helpers for WebUI image-generation intent metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
IMAGE_GENERATION_METADATA_KEY = "image_generation"
|
||||
|
||||
|
||||
def image_generation_prompt(content: str, metadata: dict[str, Any] | None) -> str:
|
||||
"""Decorate a user prompt when WebUI image mode is enabled."""
|
||||
runtime_context = image_generation_runtime_context(metadata)
|
||||
return f"{content}\n\n{runtime_context}" if runtime_context else content
|
||||
|
||||
|
||||
def image_generation_runtime_context(metadata: dict[str, Any] | None) -> str:
|
||||
"""Return the model-only instruction for WebUI image generation mode."""
|
||||
raw = (metadata or {}).get(IMAGE_GENERATION_METADATA_KEY)
|
||||
if not isinstance(raw, dict) or raw.get("enabled") is not True:
|
||||
return ""
|
||||
|
||||
aspect_ratio = raw.get("aspect_ratio")
|
||||
if isinstance(aspect_ratio, str) and aspect_ratio.strip():
|
||||
instruction = (
|
||||
"The user selected WebUI image generation mode. Use the generate_image tool. "
|
||||
f"When calling generate_image, pass aspect_ratio={aspect_ratio!r}."
|
||||
)
|
||||
else:
|
||||
instruction = (
|
||||
"The user selected WebUI image generation mode. Use the generate_image tool. "
|
||||
"Choose the most suitable aspect_ratio yourself from the prompt and intended use."
|
||||
)
|
||||
return f"[WebUI image generation instruction: {instruction}]"
|
||||
@@ -11,9 +11,8 @@ 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 LLMProvider, LLMResponse, ToolCallRequest
|
||||
from nanobot.providers.base import LLMResponse, ToolCallRequest
|
||||
from nanobot.providers.image_generation import GeneratedImageResponse
|
||||
from nanobot.runtime_context import public_history_message
|
||||
|
||||
PNG_DATA_URL = (
|
||||
"data:image/png;base64,"
|
||||
@@ -86,56 +85,3 @@ async def test_outbound_no_longer_carries_generated_media(
|
||||
# OutboundMessage no longer carries generated media —
|
||||
# the LLM sends images via the message tool instead.
|
||||
assert result.media == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_mode_instruction_is_persisted_as_next_turn_prefix(tmp_path: Path) -> None:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.generation.max_tokens = 4096
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="first answer"),
|
||||
LLMResponse(content="second answer"),
|
||||
])
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
tools_config=ToolsConfig(
|
||||
image_generation=ImageGenerationToolConfig(enabled=True),
|
||||
),
|
||||
)
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False)
|
||||
|
||||
await loop._process_message(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="chat-image-prefix",
|
||||
content="draw a fox",
|
||||
metadata={
|
||||
"image_generation": {
|
||||
"enabled": True,
|
||||
"aspect_ratio": "16:9",
|
||||
},
|
||||
},
|
||||
))
|
||||
await loop._process_message(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="chat-image-prefix",
|
||||
content="thanks",
|
||||
))
|
||||
|
||||
first_wire = LLMProvider._sanitize_empty_content(
|
||||
provider.chat_with_retry.await_args_list[0].kwargs["messages"]
|
||||
)
|
||||
second_wire = LLMProvider._sanitize_empty_content(
|
||||
provider.chat_with_retry.await_args_list[1].kwargs["messages"]
|
||||
)
|
||||
assert second_wire[: len(first_wire)] == first_wire
|
||||
assert "aspect_ratio='16:9'" in first_wire[1]["content"]
|
||||
|
||||
persisted = loop.sessions.get_or_create("websocket:chat-image-prefix").messages[0]
|
||||
assert persisted["content"] == first_wire[1]["content"]
|
||||
assert public_history_message(persisted)["content"] == "draw a fox"
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
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
|
||||
@@ -13,7 +13,6 @@ import type { StreamError } from "@/lib/nanobot-client";
|
||||
import type {
|
||||
InboundEvent,
|
||||
OutboundCliAppMention,
|
||||
OutboundImageGeneration,
|
||||
OutboundMcpPresetMention,
|
||||
OutboundMedia,
|
||||
GoalStateWsPayload,
|
||||
@@ -477,7 +476,6 @@ export interface SendImage {
|
||||
}
|
||||
|
||||
export interface SendOptions {
|
||||
imageGeneration?: OutboundImageGeneration;
|
||||
cliApps?: OutboundCliAppMention[];
|
||||
mcpPresets?: OutboundMcpPresetMention[];
|
||||
workspaceScope?: WorkspaceScopePayload | null;
|
||||
|
||||
@@ -3,7 +3,6 @@ import type {
|
||||
InboundEvent,
|
||||
Outbound,
|
||||
OutboundCliAppMention,
|
||||
OutboundImageGeneration,
|
||||
OutboundMcpPresetMention,
|
||||
OutboundMedia,
|
||||
GoalStateWsPayload,
|
||||
@@ -385,7 +384,6 @@ export class NanobotClient {
|
||||
content: string,
|
||||
media?: OutboundMedia[],
|
||||
options?: {
|
||||
imageGeneration?: OutboundImageGeneration;
|
||||
cliApps?: OutboundCliAppMention[];
|
||||
mcpPresets?: OutboundMcpPresetMention[];
|
||||
workspaceScope?: WorkspaceScopePayload | null;
|
||||
@@ -398,7 +396,6 @@ export class NanobotClient {
|
||||
chat_id: chatId,
|
||||
content,
|
||||
...(media && media.length > 0 ? { media } : {}),
|
||||
...(options?.imageGeneration ? { image_generation: options.imageGeneration } : {}),
|
||||
...(options?.cliApps?.length ? { cli_apps: options.cliApps } : {}),
|
||||
...(options?.mcpPresets?.length ? { mcp_presets: options.mcpPresets } : {}),
|
||||
...(options?.workspaceScope ? { workspace_scope: options.workspaceScope } : {}),
|
||||
|
||||
@@ -932,11 +932,6 @@ export interface OutboundMedia {
|
||||
name?: string;
|
||||
}
|
||||
|
||||
export interface OutboundImageGeneration {
|
||||
enabled: true;
|
||||
aspect_ratio?: string | null;
|
||||
}
|
||||
|
||||
export interface OutboundCliAppMention {
|
||||
name: string;
|
||||
display_name?: string;
|
||||
@@ -998,7 +993,6 @@ export type Outbound =
|
||||
chat_id: string;
|
||||
content: string;
|
||||
media?: OutboundMedia[];
|
||||
image_generation?: OutboundImageGeneration;
|
||||
cli_apps?: OutboundCliAppMention[];
|
||||
mcp_presets?: OutboundMcpPresetMention[];
|
||||
workspace_scope?: WorkspaceScopePayload;
|
||||
|
||||
@@ -528,33 +528,6 @@ describe("NanobotClient", () => {
|
||||
});
|
||||
});
|
||||
|
||||
it("includes image generation options in outbound messages", () => {
|
||||
const client = new NanobotClient({
|
||||
url: "ws://test",
|
||||
reconnect: false,
|
||||
socketFactory: (url) => new FakeSocket(url) as unknown as WebSocket,
|
||||
});
|
||||
client.connect();
|
||||
lastSocket().fakeOpen();
|
||||
|
||||
client.sendMessage(
|
||||
"chat-img",
|
||||
"draw a banner",
|
||||
undefined,
|
||||
{ imageGeneration: { enabled: true, aspect_ratio: "16:9" } },
|
||||
);
|
||||
|
||||
expect(lastSocket().sent).toContain(
|
||||
JSON.stringify({
|
||||
type: "message",
|
||||
chat_id: "chat-img",
|
||||
content: "draw a banner",
|
||||
image_generation: { enabled: true, aspect_ratio: "16:9" },
|
||||
webui: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("includes CLI app attachments in outbound messages", () => {
|
||||
const client = new NanobotClient({
|
||||
url: "ws://test",
|
||||
|
||||
@@ -1640,31 +1640,6 @@ describe("useNanobotStream", () => {
|
||||
expect(result.current.messages[0].media).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("passes image generation options to the websocket client", () => {
|
||||
const fake = fakeClient();
|
||||
const { result } = renderHook(() => useNanobotStream("chat-img", EMPTY_MESSAGES), {
|
||||
wrapper: wrap(fake.client),
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.send(
|
||||
"draw a square icon",
|
||||
undefined,
|
||||
{ imageGeneration: { enabled: true, aspect_ratio: "1:1" } },
|
||||
);
|
||||
});
|
||||
|
||||
expect(fake.client.sendMessage).toHaveBeenCalledWith(
|
||||
"chat-img",
|
||||
"draw a square icon",
|
||||
undefined,
|
||||
expect.objectContaining({
|
||||
imageGeneration: { enabled: true, aspect_ratio: "1:1" },
|
||||
turnId: expect.any(String),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("stops the active turn without adding a user slash command bubble", () => {
|
||||
const fake = fakeClient();
|
||||
const { result } = renderHook(() => useNanobotStream("chat-stop", EMPTY_MESSAGES), {
|
||||
|
||||
Reference in New Issue
Block a user