diff --git a/packages/agent/test/e2e.test.ts b/packages/agent/test/e2e.test.ts index 052d327b..f032555e 100644 --- a/packages/agent/test/e2e.test.ts +++ b/packages/agent/test/e2e.test.ts @@ -7,7 +7,6 @@ import { fauxToolCall, type Model, registerFauxProvider, - type TextContent, type ToolResultMessage, type UserMessage, } from "@mariozechner/pi-ai"; @@ -24,8 +23,10 @@ function createFauxRegistration(options: Parameters } function getTextContent(message: AssistantMessage | ToolResultMessage): string { - const textBlocks = message.content.filter((block): block is TextContent => block.type === "text") as TextContent[]; - return textBlocks.map((block) => block.text).join("\n"); + return message.content + .filter((block) => block.type === "text") + .map((block) => block.text) + .join("\n"); } afterEach(() => { diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 9f0e60f0..8050ec83 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -15,8 +15,6 @@ ### 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 diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 4b1a7d75..2c848ae4 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -88,29 +88,20 @@ function getBedrockBaseUrl(modelId: string): string { async function fetchOpenRouterModels(): Promise[]> { try { console.log("Fetching models from OpenRouter API..."); - 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(); - for (const model of [...generalData.data, ...imageData.data]) { - combinedModels.set(model.id, model); - } + const response = await fetch("https://openrouter.ai/api/v1/models"); + const data = await response.json(); const models: Model[] = []; - for (const model of combinedModels.values()) { - // OpenRouter model IDs already include the upstream provider prefix. - // Keep the full ID (e.g. "google/gemini-2.5-flash-image"). - const provider: KnownProvider = "openrouter"; - const modelKey = model.id; + for (const model of data.data) { + // Only include models that support tools + if (!model.supported_parameters?.includes("tools")) continue; - const supportsTools = model.supported_parameters?.includes("tools"); - const supportsImageOutput = model.architecture?.output_modalities?.includes("image"); - if (!supportsTools && !supportsImageOutput) continue; + // Parse provider from model ID + let provider: KnownProvider = "openrouter"; + let modelKey = model.id; + + modelKey = model.id; // Keep full ID for OpenRouter // Parse input modalities const input: ("text" | "image")[] = ["text"]; @@ -132,7 +123,6 @@ async function fetchOpenRouterModels(): Promise[]> { provider, reasoning: model.supported_parameters?.includes("reasoning") || false, input, - compat: supportsImageOutput ? { openRouterImageGeneration: true } : undefined, cost: { input: inputCost, output: outputCost, @@ -140,12 +130,12 @@ async function fetchOpenRouterModels(): Promise[]> { cacheWrite: cacheWriteCost, }, contextWindow: model.context_length || 4096, - maxTokens: model.top_provider?.max_completion_tokens || model.context_length || 4096, + maxTokens: model.top_provider?.max_completion_tokens || 4096, }; models.push(normalizedModel); } - console.log(`Fetched ${models.length} tool-capable or image-generation models from OpenRouter`); + console.log(`Fetched ${models.length} tool-capable models from OpenRouter`); return models; } catch (error) { console.error("Failed to fetch OpenRouter models:", error); diff --git a/packages/ai/src/providers/faux.ts b/packages/ai/src/providers/faux.ts index 8c0277b0..7bba1b72 100644 --- a/packages/ai/src/providers/faux.ts +++ b/packages/ai/src/providers/faux.ts @@ -147,7 +147,7 @@ function contentToText(content: string | Array): str .join("\n"); } -function assistantContentToText(content: Array): string { +function assistantContentToText(content: Array): string { return content .map((block) => { if (block.type === "text") { @@ -156,9 +156,6 @@ function assistantContentToText(content: Array } } - 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") { diff --git a/packages/ai/src/providers/mistral.ts b/packages/ai/src/providers/mistral.ts index 501b1f53..68affc98 100644 --- a/packages/ai/src/providers/mistral.ts +++ b/packages/ai/src/providers/mistral.ts @@ -524,9 +524,6 @@ function toChatMessages(messages: Message[], supportsImages: boolean): ChatCompl } continue; } - if (block.type !== "toolCall") { - continue; - } toolCalls.push({ id: block.id, type: "function", diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index f2d36e7c..f0e167e9 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -1,6 +1,5 @@ import OpenAI from "openai"; import type { - ChatCompletion, ChatCompletionAssistantMessageParam, ChatCompletionChunk, ChatCompletionContentPart, @@ -99,26 +98,6 @@ 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; @@ -171,75 +150,6 @@ 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(); @@ -575,10 +485,6 @@ function buildParams( prompt_cache_retention: cacheRetention === "long" && compat.supportsLongCacheRetention ? "24h" : undefined, }; - if (isOpenRouterImageGenerationModel(model)) { - Reflect.set(params, "modalities", ["image"]); - } - if (compat.supportsUsageInStreaming !== false) { (params as any).stream_options = { include_usage: true }; } @@ -599,12 +505,12 @@ function buildParams( params.temperature = options.temperature; } - if (!isOpenRouterImageGenerationModel(model) && context.tools && context.tools.length > 0) { + if (context.tools && context.tools.length > 0) { params.tools = convertTools(context.tools, compat); if (compat.zaiToolStream) { (params as any).tool_stream = true; } - } else if (!isOpenRouterImageGenerationModel(model) && hasToolHistory(context.messages)) { + } else if (hasToolHistory(context.messages)) { // Anthropic (via LiteLLM/proxy) requires tools param when conversation has tool_calls/tool_results params.tools = []; } @@ -1174,7 +1080,6 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet ? "openrouter" : "openai", openRouterRouting: {}, - openRouterImageGeneration: false, vercelGatewayRouting: {}, zaiToolStream: false, supportsStrictMode: true, @@ -1208,7 +1113,6 @@ 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, diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 9a4c7a9f..8ffd785f 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -212,7 +212,7 @@ export interface UserMessage { export interface AssistantMessage { role: "assistant"; - content: (TextContent | ThinkingContent | ImageContent | ToolCall)[]; + content: (TextContent | ThinkingContent | ToolCall)[]; api: Api; provider: Provider; model: string; @@ -265,8 +265,6 @@ 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 } @@ -302,8 +300,6 @@ 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. */ diff --git a/packages/ai/test/google-image-output.test.ts b/packages/ai/test/google-image-output.test.ts deleted file mode 100644 index 13563093..00000000 --- a/packages/ai/test/google-image-output.test.ts +++ /dev/null @@ -1,103 +0,0 @@ -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=", - }); - }); -}); diff --git a/packages/ai/test/openai-completions-image-output.test.ts b/packages/ai/test/openai-completions-image-output.test.ts deleted file mode 100644 index 745305c4..00000000 --- a/packages/ai/test/openai-completions-image-output.test.ts +++ /dev/null @@ -1,100 +0,0 @@ -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 & { - 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", - compat: { openRouterImageGeneration: true }, - 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; modalities?: string[] }; - expect(params.stream).toBe(false); - expect("stream_options" in (params as object)).toBe(false); - expect(params.modalities).toEqual(["image"]); - }); -}); diff --git a/packages/ai/test/openai-completions-thinking-as-text.test.ts b/packages/ai/test/openai-completions-thinking-as-text.test.ts index 012e434b..24425f3a 100644 --- a/packages/ai/test/openai-completions-thinking-as-text.test.ts +++ b/packages/ai/test/openai-completions-thinking-as-text.test.ts @@ -34,7 +34,6 @@ const compat = { requiresReasoningContentOnAssistantMessages: false, thinkingFormat: "openai", openRouterRouting: {}, - openRouterImageGeneration: false, vercelGatewayRouting: {}, zaiToolStream: false, supportsStrictMode: true, diff --git a/packages/ai/test/openai-completions-tool-result-images.test.ts b/packages/ai/test/openai-completions-tool-result-images.test.ts index 8154c9ac..ac1eb28a 100644 --- a/packages/ai/test/openai-completions-tool-result-images.test.ts +++ b/packages/ai/test/openai-completions-tool-result-images.test.ts @@ -32,7 +32,6 @@ const compat: Required = { requiresReasoningContentOnAssistantMessages: false, thinkingFormat: "openai", openRouterRouting: {}, - openRouterImageGeneration: false, vercelGatewayRouting: {}, zaiToolStream: false, supportsStrictMode: true, diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 26299b2b..f366274b 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -502,8 +502,10 @@ function extractTextContent(message: Message): string { if (typeof content === "string") { return content; } - const textBlocks = content.filter((block): block is TextContent => block.type === "text") as TextContent[]; - return textBlocks.map((block) => block.text).join(" "); + return content + .filter((block): block is TextContent => block.type === "text") + .map((block) => block.text) + .join(" "); } function getLastActivityTime(entries: FileEntry[]): number | undefined { diff --git a/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts b/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts index 678d6400..626365c6 100644 --- a/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts +++ b/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts @@ -1,7 +1,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { fauxAssistantMessage, registerFauxProvider, type TextContent } from "@mariozechner/pi-ai"; +import { fauxAssistantMessage, registerFauxProvider } from "@mariozechner/pi-ai"; import { afterEach, describe, expect, it } from "vitest"; import type { AgentSession } from "../../../src/core/agent-session.js"; import { @@ -18,11 +18,12 @@ function getText(message: AgentSession["messages"][number]): string { if (!("content" in message)) { return ""; } - 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(""); + 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(""); } describe("regression #2860: replaced session callbacks", () => {