shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: image content

Browse Source
This commit is contained in:
Cristina Poncela Cubeiro
2026-04-28 13:43:28 +02:00
parent 9848b3145a
commit c3c10737d8
14 changed files with 385 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
packages/agent/test/e2e.test.ts
View File

@@ -7,6 +7,7 @@ import {
fauxToolCall,
type Model,
registerFauxProvider,
type TextContent,
type ToolResultMessage,
type UserMessage,
} from "@mariozechner/pi-ai";
@@ -23,10 +24,8 @@ function createFauxRegistration(options: Parameters<typeof registerFauxProvider>
}
function getTextContent(message: AssistantMessage | ToolResultMessage): string {
return message.content
.filter((block) => block.type === "text")
.map((block) => block.text)
.join("\n");
const textBlocks = message.content.filter((block): block is TextContent => block.type === "text") as TextContent[];
return textBlocks.map((block) => block.text).join("\n");
}
afterEach(() => {

2
packages/ai/CHANGELOG.md
View File

@@ -15,6 +15,8 @@
### Added
- Added Azure Cognitive Services endpoint support for Azure OpenAI Responses base URLs ([#3799](https://github.com/badlogic/pi-mono/pull/3799) by [@marcbloech](https://github.com/marcbloech)).
- Added OpenRouter image-generation support by registering Flux and Gemini image-preview models, accepting OpenRouter image-output models from discovery, and parsing generated image payloads from OpenAI-compatible completions responses.
- Added Google streaming to surface inline image output parts as assistant image content blocks and `image_start`/`image_end` events.
### Changed

26
packages/ai/scripts/generate-models.ts
View File

@@ -88,20 +88,29 @@ function getBedrockBaseUrl(modelId: string): string {
async function fetchOpenRouterModels(): Promise<Model<any>[]> {
try {
console.log("Fetching models from OpenRouter API...");
const response = await fetch("https://openrouter.ai/api/v1/models");
const data = await response.json();
const [generalResponse, imageResponse] = await Promise.all([
fetch("https://openrouter.ai/api/v1/models"),
fetch("https://openrouter.ai/api/v1/models?output_modalities=image"),
]);
const generalData = await generalResponse.json();
const imageData = await imageResponse.json();
const combinedModels = new Map<string, any>();
for (const model of [...generalData.data, ...imageData.data]) {
combinedModels.set(model.id, model);
}
const models: Model<any>[] = [];
for (const model of data.data) {
// Only include models that support tools
if (!model.supported_parameters?.includes("tools")) continue;
// Parse provider from model ID
let provider: KnownProvider = "openrouter";
let modelKey = model.id;
modelKey = model.id; // Keep full ID for OpenRouter
for (const model of combinedModels.values()) {
const supportsTools = model.supported_parameters?.includes("tools");
const supportsImageOutput = model.architecture?.output_modalities?.includes("image");
if (!supportsTools && !supportsImageOutput) continue;
// Parse input modalities
const input: ("text" | "image")[] = ["text"];
@@ -123,6 +132,7 @@ async function fetchOpenRouterModels(): Promise<Model<any>[]> {
provider,
reasoning: model.supported_parameters?.includes("reasoning") || false,
input,
compat: supportsImageOutput ? { openRouterImageGeneration: true } : undefined,
cost: {
input: inputCost,
output: outputCost,
@@ -130,12 +140,12 @@ async function fetchOpenRouterModels(): Promise<Model<any>[]> {
cacheWrite: cacheWriteCost,
},
contextWindow: model.context_length || 4096,
maxTokens: model.top_provider?.max_completion_tokens || 4096,
maxTokens: model.top_provider?.max_completion_tokens || model.context_length || 4096,
};
models.push(normalizedModel);
}
console.log(`Fetched ${models.length} tool-capable models from OpenRouter`);
console.log(`Fetched ${models.length} tool-capable or image-generation models from OpenRouter`);
return models;
} catch (error) {
console.error("Failed to fetch OpenRouter models:", error);

12
packages/ai/src/providers/faux.ts
View File

@@ -147,7 +147,7 @@ function contentToText(content: string | Array<TextContent | ImageContent>): str
.join("\n");
}
function assistantContentToText(content: Array<TextContent | ThinkingContent | ToolCall>): string {
function assistantContentToText(content: Array<TextContent | ThinkingContent | ImageContent | ToolCall>): string {
return content
.map((block) => {
if (block.type === "text") {
@@ -156,6 +156,9 @@ function assistantContentToText(content: Array<TextContent | ThinkingContent | T
if (block.type === "thinking") {
return block.thinking;
}
if (block.type === "image") {
return `[image:${block.mimeType}]`;
}
return `${block.name}:${JSON.stringify(block.arguments)}`;
})
.join("\n");
@@ -362,6 +365,13 @@ async function streamWithDeltas(
continue;
}
if (block.type === "image") {
partial.content = [...partial.content, block];
stream.push({ type: "image_start", contentIndex: index, partial: { ...partial } });
stream.push({ type: "image_end", contentIndex: index, image: block, partial: { ...partial } });
continue;
}
partial.content = [...partial.content, { type: "toolCall", id: block.id, name: block.name, arguments: {} }];
stream.push({ type: "toolcall_start", contentIndex: index, partial: { ...partial } });
for (const chunk of splitStringByTokenSize(JSON.stringify(block.arguments), minTokenSize, maxTokenSize)) {

38
packages/ai/src/providers/google.ts
View File

@@ -10,6 +10,7 @@ import type {
Api,
AssistantMessage,
Context,
ImageContent,
Model,
SimpleStreamOptions,
StreamFunction,
@@ -153,6 +154,43 @@ export const streamGoogle: StreamFunction<"google-generative-ai", GoogleOptions>
}
}
if (part.inlineData) {
if (currentBlock) {
if (currentBlock.type === "text") {
stream.push({
type: "text_end",
contentIndex: blockIndex(),
content: currentBlock.text,
partial: output,
});
} else {
stream.push({
type: "thinking_end",
contentIndex: blockIndex(),
content: currentBlock.thinking,
partial: output,
});
}
currentBlock = null;
}
if (part.inlineData.data && part.inlineData.mimeType) {
const imageBlock: ImageContent = {
type: "image",
data: part.inlineData.data,
mimeType: part.inlineData.mimeType,
};
output.content.push(imageBlock);
stream.push({ type: "image_start", contentIndex: blockIndex(), partial: output });
stream.push({
type: "image_end",
contentIndex: blockIndex(),
image: imageBlock,
partial: output,
});
}
}
if (part.functionCall) {
if (currentBlock) {
if (currentBlock.type === "text") {

3
packages/ai/src/providers/mistral.ts
View File

@@ -524,6 +524,9 @@ function toChatMessages(messages: Message[], supportsImages: boolean): ChatCompl
}
continue;
}
if (block.type !== "toolCall") {
continue;
}
toolCalls.push({
id: block.id,
type: "function",

96
packages/ai/src/providers/openai-completions.ts
View File

@@ -1,5 +1,6 @@
import OpenAI from "openai";
import type {
ChatCompletion,
ChatCompletionAssistantMessageParam,
ChatCompletionChunk,
ChatCompletionContentPart,
@@ -98,6 +99,26 @@ type ChatCompletionToolWithCacheControl = OpenAI.Chat.Completions.ChatCompletion
cache_control?: OpenAICompatCacheControl;
};
interface OpenRouterGeneratedImage {
image_url?: string | { url?: string };
}
type OpenRouterImageGenerationMessage = ChatCompletion["choices"][number]["message"] & {
images?: OpenRouterGeneratedImage[];
};
type OpenRouterImageGenerationChoice = ChatCompletion["choices"][number] & {
message: OpenRouterImageGenerationMessage;
};
type OpenRouterImageGenerationResponse = ChatCompletion & {
choices: OpenRouterImageGenerationChoice[];
};
function isOpenRouterImageGenerationModel(model: Model<"openai-completions">): boolean {
return model.provider === "openrouter" && model.compat?.openRouterImageGeneration === true;
}
function resolveCacheRetention(cacheRetention?: CacheRetention): CacheRetention {
if (cacheRetention) {
return cacheRetention;
@@ -150,6 +171,75 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
...(options?.timeoutMs !== undefined ? { timeout: options.timeoutMs } : {}),
...(options?.maxRetries !== undefined ? { maxRetries: options.maxRetries } : {}),
};
if (isOpenRouterImageGenerationModel(model)) {
const { stream_options: _streamOptions, ...nonStreamingBaseParams } = params;
const nonStreamingParams = {
...nonStreamingBaseParams,
stream: false,
} as OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming;
const { data: response, response: rawResponse } = await client.chat.completions
.create(nonStreamingParams, requestOptions)
.withResponse();
await options?.onResponse?.(
{ status: rawResponse.status, headers: headersToRecord(rawResponse.headers) },
model,
);
stream.push({ type: "start", partial: output });
const imageResponse = response as OpenRouterImageGenerationResponse;
output.responseId ||= imageResponse.id;
output.usage = parseChunkUsage(imageResponse.usage ?? {}, model);
const choice = imageResponse.choices[0];
if (choice) {
const text = choice.message.content;
if (text) {
const textBlock: TextContent = { type: "text", text };
output.content.push(textBlock);
const contentIndex = output.content.length - 1;
stream.push({ type: "text_start", contentIndex, partial: output });
stream.push({ type: "text_delta", contentIndex, delta: text, partial: output });
stream.push({ type: "text_end", contentIndex, content: text, partial: output });
}
for (const image of choice.message.images ?? []) {
const imageUrl = typeof image.image_url === "string" ? image.image_url : image.image_url?.url;
if (!imageUrl?.startsWith("data:")) {
continue;
}
const matches = imageUrl.match(/^data:([^;]+);base64,(.+)$/);
if (!matches) {
continue;
}
const imageBlock: ImageContent = {
type: "image",
mimeType: matches[1],
data: matches[2],
};
output.content.push(imageBlock);
const contentIndex = output.content.length - 1;
stream.push({ type: "image_start", contentIndex, partial: output });
stream.push({ type: "image_end", contentIndex, image: imageBlock, partial: output });
}
if (choice.finish_reason) {
const finishReasonResult = mapStopReason(choice.finish_reason);
output.stopReason = finishReasonResult.stopReason;
if (finishReasonResult.errorMessage) {
output.errorMessage = finishReasonResult.errorMessage;
}
}
}
if (output.stopReason === "stop" || output.stopReason === "length" || output.stopReason === "toolUse") {
stream.push({ type: "done", reason: output.stopReason, message: output });
} else {
stream.push({ type: "error", reason: output.stopReason, error: output });
}
stream.end();
return;
}
const { data: openaiStream, response } = await client.chat.completions
.create(params, requestOptions)
.withResponse();
@@ -505,12 +595,12 @@ function buildParams(
params.temperature = options.temperature;
}
if (context.tools && context.tools.length > 0) {
if (!isOpenRouterImageGenerationModel(model) && context.tools && context.tools.length > 0) {
params.tools = convertTools(context.tools, compat);
if (compat.zaiToolStream) {
(params as any).tool_stream = true;
}
} else if (hasToolHistory(context.messages)) {
} else if (!isOpenRouterImageGenerationModel(model) && hasToolHistory(context.messages)) {
// Anthropic (via LiteLLM/proxy) requires tools param when conversation has tool_calls/tool_results
params.tools = [];
}
@@ -1080,6 +1170,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
? "openrouter"
: "openai",
openRouterRouting: {},
openRouterImageGeneration: false,
vercelGatewayRouting: {},
zaiToolStream: false,
supportsStrictMode: true,
@@ -1113,6 +1204,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
detected.requiresReasoningContentOnAssistantMessages,
thinkingFormat: model.compat.thinkingFormat ?? detected.thinkingFormat,
openRouterRouting: model.compat.openRouterRouting ?? {},
openRouterImageGeneration: model.compat.openRouterImageGeneration ?? detected.openRouterImageGeneration,
vercelGatewayRouting: model.compat.vercelGatewayRouting ?? detected.vercelGatewayRouting,
zaiToolStream: model.compat.zaiToolStream ?? detected.zaiToolStream,
supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode,

6
packages/ai/src/types.ts
View File

@@ -212,7 +212,7 @@ export interface UserMessage {
export interface AssistantMessage {
role: "assistant";
content: (TextContent | ThinkingContent | ToolCall)[];
content: (TextContent | ThinkingContent | ImageContent | ToolCall)[];
api: Api;
provider: Provider;
model: string;
@@ -265,6 +265,8 @@ export type AssistantMessageEvent =
| { type: "thinking_start"; contentIndex: number; partial: AssistantMessage }
| { type: "thinking_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
| { type: "thinking_end"; contentIndex: number; content: string; partial: AssistantMessage }
| { type: "image_start"; contentIndex: number; partial: AssistantMessage }
| { type: "image_end"; contentIndex: number; image: ImageContent; partial: AssistantMessage }
| { type: "toolcall_start"; contentIndex: number; partial: AssistantMessage }
| { type: "toolcall_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
| { type: "toolcall_end"; contentIndex: number; toolCall: ToolCall; partial: AssistantMessage }
@@ -300,6 +302,8 @@ export interface OpenAICompletionsCompat {
thinkingFormat?: "openai" | "openrouter" | "deepseek" | "zai" | "qwen" | "qwen-chat-template";
/** OpenRouter-specific routing preferences. Only used when baseUrl points to OpenRouter. */
openRouterRouting?: OpenRouterRouting;
/** Whether the model uses OpenRouter's image-generation response shape with assistant `message.images`. */
openRouterImageGeneration?: boolean;
/** Vercel AI Gateway routing preferences. Only used when baseUrl points to Vercel AI Gateway. */
vercelGatewayRouting?: VercelGatewayRouting;
/** Whether z.ai supports top-level `tool_stream: true` for streaming tool call deltas. Default: false. */

103
packages/ai/test/google-image-output.test.ts Normal file
View File

@@ -0,0 +1,103 @@
import { describe, expect, it, vi } from "vitest";
vi.mock("@google/genai", () => {
class GoogleGenAI {
models = {
generateContentStream: async function* () {
yield {
responseId: "google-image-1",
candidates: [
{
content: {
parts: [
{ text: "Here is your image." },
{ inlineData: { mimeType: "image/png", data: "ZmFrZS1wbmc=" } },
],
},
finishReason: "STOP",
},
],
usageMetadata: {
promptTokenCount: 5,
candidatesTokenCount: 7,
totalTokenCount: 12,
},
};
},
};
}
return {
GoogleGenAI,
FinishReason: {
STOP: "STOP",
MAX_TOKENS: "MAX_TOKENS",
BLOCKLIST: "BLOCKLIST",
PROHIBITED_CONTENT: "PROHIBITED_CONTENT",
SPII: "SPII",
SAFETY: "SAFETY",
IMAGE_SAFETY: "IMAGE_SAFETY",
IMAGE_PROHIBITED_CONTENT: "IMAGE_PROHIBITED_CONTENT",
IMAGE_RECITATION: "IMAGE_RECITATION",
IMAGE_OTHER: "IMAGE_OTHER",
RECITATION: "RECITATION",
FINISH_REASON_UNSPECIFIED: "FINISH_REASON_UNSPECIFIED",
OTHER: "OTHER",
LANGUAGE: "LANGUAGE",
MALFORMED_FUNCTION_CALL: "MALFORMED_FUNCTION_CALL",
UNEXPECTED_TOOL_CALL: "UNEXPECTED_TOOL_CALL",
NO_IMAGE: "NO_IMAGE",
},
FunctionCallingConfigMode: {
AUTO: "AUTO",
ANY: "ANY",
NONE: "NONE",
},
ThinkingLevel: {
THINKING_LEVEL_UNSPECIFIED: "THINKING_LEVEL_UNSPECIFIED",
MINIMAL: "MINIMAL",
LOW: "LOW",
MEDIUM: "MEDIUM",
HIGH: "HIGH",
},
};
});
import { streamGoogle } from "../src/providers/google.js";
import type { Context, Model } from "../src/types.js";
describe("google image output", () => {
it("emits assistant image blocks for inlineData parts", async () => {
const model: Model<"google-generative-ai"> = {
id: "gemini-2.5-flash",
name: "Gemini 2.5 Flash",
api: "google-generative-ai",
provider: "google",
baseUrl: "https://generativelanguage.googleapis.com",
reasoning: false,
input: ["text", "image"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 1000000,
maxTokens: 8192,
};
const context: Context = {
messages: [{ role: "user", content: "Generate an image", timestamp: Date.now() }],
};
const result = streamGoogle(model, context, { apiKey: "test" });
const eventTypes: string[] = [];
for await (const event of result) {
eventTypes.push(event.type);
}
const message = await result.result();
expect(eventTypes).toEqual(["start", "text_start", "text_delta", "text_end", "image_start", "image_end", "done"]);
expect(message.responseId).toBe("google-image-1");
expect(message.content[0]).toMatchObject({ type: "text", text: "Here is your image." });
expect(message.content[1]).toMatchObject({
type: "image",
mimeType: "image/png",
data: "ZmFrZS1wbmc=",
});
});
});

98
packages/ai/test/openai-completions-image-output.test.ts Normal file
View File

@@ -0,0 +1,98 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import { streamSimple } from "../src/stream.js";
import type { Context, Model } from "../src/types.js";
const mockState = vi.hoisted(() => ({
lastParams: undefined as unknown,
}));
vi.mock("openai", () => {
class FakeOpenAI {
chat = {
completions: {
create: (params: unknown) => {
mockState.lastParams = params;
const response = {
id: "chatcmpl-image-1",
usage: {
prompt_tokens: 12,
completion_tokens: 34,
prompt_tokens_details: { cached_tokens: 0 },
},
choices: [
{
finish_reason: "stop",
message: {
role: "assistant",
content: "Here is your image.",
images: [
{
image_url: "data:image/png;base64,ZmFrZS1wbmc=",
},
],
},
},
],
};
const promise = Promise.resolve(response) as Promise<typeof response> & {
withResponse: () => Promise<{
data: typeof response;
response: { status: number; headers: Headers };
}>;
};
promise.withResponse = async () => ({
data: response,
response: { status: 200, headers: new Headers() },
});
return promise;
},
},
};
}
return { default: FakeOpenAI };
});
describe("openai-completions image output", () => {
beforeEach(() => {
mockState.lastParams = undefined;
});
it("switches OpenRouter image generation models to non-streaming and emits image events", async () => {
const model: Model<"openai-completions"> = {
id: "black-forest-labs/flux.2-pro",
name: "Black Forest Labs: FLUX.2 Pro",
api: "openai-completions",
provider: "openrouter",
baseUrl: "https://openrouter.ai/api/v1",
reasoning: false,
input: ["text", "image"],
cost: { input: 0.015, output: 0.03, cacheRead: 0, cacheWrite: 0 },
contextWindow: 100000,
maxTokens: 100000,
};
const context: Context = {
messages: [{ role: "user", content: "Generate a cat", timestamp: Date.now() }],
};
const result = streamSimple(model, context, { apiKey: "test" });
const eventTypes: string[] = [];
for await (const event of result) {
eventTypes.push(event.type);
}
const message = await result.result();
expect(eventTypes).toEqual(["start", "text_start", "text_delta", "text_end", "image_start", "image_end", "done"]);
expect(message.responseId).toBe("chatcmpl-image-1");
expect(message.content[0]).toMatchObject({ type: "text", text: "Here is your image." });
expect(message.content[1]).toMatchObject({
type: "image",
mimeType: "image/png",
data: "ZmFrZS1wbmc=",
});
const params = mockState.lastParams as { stream?: boolean; stream_options?: unknown };
expect(params.stream).toBe(false);
expect("stream_options" in (params as object)).toBe(false);
});
});

1
packages/ai/test/openai-completions-thinking-as-text.test.ts
View File

@@ -34,6 +34,7 @@ const compat = {
requiresReasoningContentOnAssistantMessages: false,
thinkingFormat: "openai",
openRouterRouting: {},
openRouterImageGeneration: false,
vercelGatewayRouting: {},
zaiToolStream: false,
supportsStrictMode: true,

1
packages/ai/test/openai-completions-tool-result-images.test.ts
View File

@@ -32,6 +32,7 @@ const compat: Required<OpenAICompletionsCompat> = {
requiresReasoningContentOnAssistantMessages: false,
thinkingFormat: "openai",
openRouterRouting: {},
openRouterImageGeneration: false,
vercelGatewayRouting: {},
zaiToolStream: false,
supportsStrictMode: true,

6
packages/coding-agent/src/core/session-manager.ts
View File

@@ -502,10 +502,8 @@ function extractTextContent(message: Message): string {
if (typeof content === "string") {
return content;
}
return content
.filter((block): block is TextContent => block.type === "text")
.map((block) => block.text)
.join(" ");
const textBlocks = content.filter((block): block is TextContent => block.type === "text") as TextContent[];
return textBlocks.map((block) => block.text).join(" ");
}
function getLastActivityTime(entries: FileEntry[]): number | undefined {

13
packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts
View File

@@ -1,7 +1,7 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fauxAssistantMessage, registerFauxProvider } from "@mariozechner/pi-ai";
import { fauxAssistantMessage, registerFauxProvider, type TextContent } from "@mariozechner/pi-ai";
import { afterEach, describe, expect, it } from "vitest";
import type { AgentSession } from "../../../src/core/agent-session.js";
import {
@@ -18,12 +18,11 @@ function getText(message: AgentSession["messages"][number]): string {
if (!("content" in message)) {
return "";
}
return typeof message.content === "string"
? message.content
: message.content
.filter((part): part is { type: "text"; text: string } => part.type === "text")
.map((part) => part.text)
.join("");
if (typeof message.content === "string") {
return message.content;
}
const textParts = message.content.filter((part): part is TextContent => part.type === "text") as TextContent[];
return textParts.map((part) => part.text).join("");
}
describe("regression #2860: replaced session callbacks", () => {