pi-openai-toolkit 0.7.2 → 0.8.0

package/src/image-generation/client.ts

@@ -0,0 +1,248 @@
+import { buildResponsesRequestHeaders } from "../responses-headers";
+import type { ResponsesRuntime } from "../runtime";
+import {
+	extractProviderErrorMessage,
+	parseImageGenerationResponse,
+	type ImageGenerationRequestBody,
+} from "./protocol";
+import {
+	IMAGE_REQUEST_TIMEOUT_MS,
+	MAX_IMAGE_ERROR_BYTES,
+	MAX_IMAGE_RESPONSE_BYTES,
+	sanitizeImageDiagnostic,
+	type ParsedGeneratedImage,
+} from "./types";
+
+export type ImageGenerationClientFailureReason =
+	| "aborted"
+	| "timeout"
+	| "authentication"
+	| "rate-limit"
+	| "request-rejected"
+	| "backend-unavailable"
+	| "network"
+	| "oversized-response"
+	| "malformed-response"
+	| "no-image";
+
+export type ImageGenerationClientResult =
+	| { ok: true; image: ParsedGeneratedImage; status: number }
+	| {
+			ok: false;
+			reason: ImageGenerationClientFailureReason;
+			status?: number;
+			errorMessage: string;
+	  };
+
+export type ImageGenerationFetch = (
+	input: string | URL | Request,
+	init?: RequestInit,
+) => Promise<Response>;
+
+function buildRequestHeaders(runtime: ResponsesRuntime): Headers {
+	return buildResponsesRequestHeaders(runtime, { accept: "application/json" });
+}
+
+async function readBoundedBody(
+	response: Response,
+	maxBytes: number,
+): Promise<{ ok: true; bytes: Uint8Array } | { ok: false }> {
+	const contentLength = response.headers.get("content-length");
+	if (contentLength !== null) {
+		const parsed = Number(contentLength);
+		if (Number.isFinite(parsed) && parsed > maxBytes) return { ok: false };
+	}
+
+	if (!response.body) return { ok: true, bytes: new Uint8Array() };
+	const reader = response.body.getReader();
+	const chunks: Uint8Array[] = [];
+	let total = 0;
+	try {
+		while (true) {
+			const next = await reader.read();
+			if (next.done) break;
+			if (!next.value) continue;
+			total += next.value.byteLength;
+			if (total > maxBytes) {
+				await reader.cancel().catch(() => undefined);
+				return { ok: false };
+			}
+			chunks.push(next.value);
+		}
+	} finally {
+		reader.releaseLock();
+	}
+
+	const bytes = new Uint8Array(total);
+	let offset = 0;
+	for (const chunk of chunks) {
+		bytes.set(chunk, offset);
+		offset += chunk.byteLength;
+	}
+	return { ok: true, bytes };
+}
+
+function parseJson(bytes: Uint8Array): unknown {
+	const text = new TextDecoder().decode(bytes);
+	return JSON.parse(text) as unknown;
+}
+
+function providerFailureText(value: unknown): string {
+	if (!value || typeof value !== "object" || Array.isArray(value)) return "";
+	const candidate = value as Record<string, unknown>;
+	const error = candidate.error && typeof candidate.error === "object" && !Array.isArray(candidate.error)
+		? candidate.error as Record<string, unknown>
+		: candidate;
+	return [error.code, error.type, error.message]
+		.filter((item): item is string => typeof item === "string")
+		.join(" ")
+		.toLowerCase();
+}
+
+function mapHttpFailure(status: number, payload?: unknown): ImageGenerationClientFailureReason {
+	if (status === 401 || status === 403) return "authentication";
+	const providerFailure = providerFailureText(payload);
+	if (
+		status === 429 ||
+		/(?:rate|usage|billing)[-_ ]?limit|quota|insufficient[-_ ]?(?:quota|credit)/.test(providerFailure)
+	) {
+		return "rate-limit";
+	}
+	if (status === 400 || status === 404 || status === 409 || status === 422) {
+		return "request-rejected";
+	}
+	if (status >= 500) return "backend-unavailable";
+	return "request-rejected";
+}
+
+function defaultHttpMessage(status: number, reason = mapHttpFailure(status)): string {
+	switch (reason) {
+		case "authentication":
+			return "Image generation authentication failed for the current model provider.";
+		case "rate-limit":
+			return "Image generation was rate-limited or the provider image quota was exhausted.";
+		case "backend-unavailable":
+			return "Image generation is temporarily unavailable at the configured gateway or upstream provider.";
+		default:
+			return `Image generation request was rejected (HTTP ${status}).`;
+	}
+}
+
+export async function requestGeneratedImage(args: {
+	runtime: ResponsesRuntime;
+	body: ImageGenerationRequestBody;
+	signal?: AbortSignal;
+	fetchFn?: ImageGenerationFetch;
+	timeoutMs?: number;
+}): Promise<ImageGenerationClientResult> {
+	const fetchFn = args.fetchFn ?? globalThis.fetch;
+	const timeoutSignal = AbortSignal.timeout(args.timeoutMs ?? IMAGE_REQUEST_TIMEOUT_MS);
+	const signal = args.signal
+		? AbortSignal.any([args.signal, timeoutSignal])
+		: timeoutSignal;
+
+	let response: Response;
+	try {
+		response = await fetchFn(args.runtime.responsesUrl, {
+			method: "POST",
+			headers: buildRequestHeaders(args.runtime),
+			body: JSON.stringify(args.body),
+			signal,
+		});
+	} catch (error) {
+		if (args.signal?.aborted) {
+			return { ok: false, reason: "aborted", errorMessage: "Image generation was cancelled." };
+		}
+		if (timeoutSignal.aborted) {
+			return {
+				ok: false,
+				reason: "timeout",
+				errorMessage: "Image generation timed out. The request was not retried automatically.",
+			};
+		}
+		return {
+			ok: false,
+			reason: "network",
+			errorMessage: sanitizeImageDiagnostic(
+				error instanceof Error ? error.message : error,
+				"Image generation network request failed. The request was not retried automatically.",
+			),
+		};
+	}
+
+	let bounded: Awaited<ReturnType<typeof readBoundedBody>>;
+	try {
+		bounded = await readBoundedBody(
+			response,
+			response.ok ? MAX_IMAGE_RESPONSE_BYTES : MAX_IMAGE_ERROR_BYTES,
+		);
+	} catch (error) {
+		if (args.signal?.aborted) {
+			return { ok: false, reason: "aborted", errorMessage: "Image generation was cancelled." };
+		}
+		if (timeoutSignal.aborted) {
+			return {
+				ok: false,
+				reason: "timeout",
+				errorMessage: "Image generation timed out. The request was not retried automatically.",
+			};
+		}
+		return {
+			ok: false,
+			reason: "network",
+			status: response.status,
+			errorMessage: sanitizeImageDiagnostic(
+				error instanceof Error ? error.message : error,
+				"Image generation response could not be read. The request was not retried automatically.",
+			),
+		};
+	}
+	if (!bounded.ok) {
+		return {
+			ok: false,
+			reason: "oversized-response",
+			status: response.status,
+			errorMessage: "Image generation response exceeded the configured size limit.",
+		};
+	}
+
+	let payload: unknown;
+	try {
+		payload = parseJson(bounded.bytes);
+	} catch {
+		return {
+			ok: false,
+			reason: "malformed-response",
+			status: response.status,
+			errorMessage: response.ok
+				? "Image generation returned invalid JSON."
+				: defaultHttpMessage(response.status),
+		};
+	}
+
+	if (!response.ok) {
+		const reason = mapHttpFailure(response.status, payload);
+		return {
+			ok: false,
+			reason,
+			status: response.status,
+			errorMessage: extractProviderErrorMessage(payload, defaultHttpMessage(response.status, reason)),
+		};
+	}
+
+	const parsed = parseImageGenerationResponse(payload);
+	if (!parsed.ok) {
+		return {
+			ok: false,
+			reason: parsed.reason,
+			status: response.status,
+			errorMessage: parsed.errorMessage,
+		};
+	}
+	return { ok: true, image: parsed.image, status: response.status };
+}
+
+export const _clientTest = {
+	buildRequestHeaders,
+	readBoundedBody,
+};

package/src/image-generation/eligibility.ts

@@ -0,0 +1,18 @@
+import { isExactModelAllowed } from "../model-scope";
+import type { ImageGenerationConfig } from "../types";
+import {
+	IMAGE_GENERATION_CAPABLE_APIS,
+	type ImageGenerationModel,
+} from "./types";
+
+export function isImageGenerationEnabledForModel(
+	model: ImageGenerationModel | undefined,
+	config: ImageGenerationConfig,
+): boolean {
+	return (
+		config.enabled &&
+		typeof model?.api === "string" &&
+		(IMAGE_GENERATION_CAPABLE_APIS as readonly string[]).includes(model.api) &&
+		isExactModelAllowed(model, config.models)
+	);
+}

package/src/image-generation/extension.ts

@@ -0,0 +1,160 @@
+import { StringEnum, Type, type Static } from "@earendil-works/pi-ai";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { loadToolkitConfig } from "../config";
+import { executeImageGeneration } from "./service";
+import { isImageGenerationEnabledForModel } from "./eligibility";
+import { renderImageGenerationResult } from "./render";
+import {
+	IMAGE_GENERATION_QUALITIES,
+	IMAGE_GENERATION_SIZES,
+	IMAGE_GENERATION_TOOL_NAME,
+	MAX_IMAGE_PATH_CHARS,
+	MAX_IMAGE_PROMPT_CHARS,
+	MAX_REFERENCE_IMAGE_COUNT,
+	ImageGenerationError,
+	sanitizeImageDiagnostic,
+	type ImageGenerationDetails,
+} from "./types";
+
+const registeredApis = new WeakSet<object>();
+
+const GenerateImageParameters = Type.Object(
+	{
+		prompt: Type.String({
+			minLength: 1,
+			maxLength: MAX_IMAGE_PROMPT_CHARS,
+			description: "Image generation or edit prompt. Preserve the user's requested subject and constraints.",
+		}),
+		referenceImagePaths: Type.Optional(
+			Type.Union(
+				[
+					Type.Null(),
+					Type.Array(Type.String({ minLength: 1, maxLength: MAX_IMAGE_PATH_CHARS }), {
+						minItems: 1,
+						maxItems: MAX_REFERENCE_IMAGE_COUNT,
+					}),
+				],
+				{
+					description:
+						"Use null when the user did not explicitly identify reference images; otherwise provide one to five user-identified local image paths.",
+				},
+			),
+		),
+		outputPath: Type.Optional(
+			Type.Union(
+				[
+					Type.Null(),
+					Type.String({
+						minLength: 1,
+						maxLength: MAX_IMAGE_PATH_CHARS,
+					}),
+				],
+				{
+					description:
+						"Use null when the user did not explicitly request a destination; otherwise provide the requested explicit .png file path.",
+				},
+			),
+		),
+		size: Type.Optional(
+			StringEnum(IMAGE_GENERATION_SIZES, {
+				description: "Requested PNG dimensions, or auto.",
+			}),
+		),
+		quality: Type.Optional(
+			StringEnum(IMAGE_GENERATION_QUALITIES, {
+				description: "Requested image quality, or auto.",
+			}),
+		),
+	},
+	{ additionalProperties: false },
+);
+
+type GenerateImageToolParams = Static<typeof GenerateImageParameters>;
+
+function syncImageGenerationTool(
+	pi: ExtensionAPI,
+	model: Parameters<typeof isImageGenerationEnabledForModel>[0],
+	config: Parameters<typeof isImageGenerationEnabledForModel>[1],
+): void {
+	const eligible = isImageGenerationEnabledForModel(model, config);
+	const activeTools = pi.getActiveTools();
+	const active = activeTools.includes(IMAGE_GENERATION_TOOL_NAME);
+	if (eligible && !active) {
+		pi.setActiveTools([...activeTools, IMAGE_GENERATION_TOOL_NAME]);
+	} else if (!eligible && active) {
+		pi.setActiveTools(activeTools.filter((name) => name !== IMAGE_GENERATION_TOOL_NAME));
+	}
+}
+
+export function registerImageGenerationExtension(
+	pi: ExtensionAPI,
+	loadConfig: typeof loadToolkitConfig = loadToolkitConfig,
+	executeImage: typeof executeImageGeneration = executeImageGeneration,
+): void {
+	if (registeredApis.has(pi)) return;
+	registeredApis.add(pi);
+
+	pi.registerTool<typeof GenerateImageParameters, ImageGenerationDetails, Record<string, never>>({
+		name: IMAGE_GENERATION_TOOL_NAME,
+		label: "OpenAI Generate Image",
+		description:
+			"Generate a PNG image, or edit from one to five user-approved local reference images, through the current allowlisted Responses model and the hosted gpt-image-2 image_generation tool. This is a paid provider operation.",
+		promptSnippet: "Generate or edit PNG images through the current allowlisted Responses model and gpt-image-2.",
+		promptGuidelines: [
+			"Use openai_generate_image when the user explicitly asks to create, draw, render, or edit a raster image.",
+			"Do not call openai_generate_image speculatively: it consumes the user's provider or gateway image quota.",
+			"Keep the image prompt faithful to the user's requested subject and constraints; do not invent unrequested style details.",
+			"Set referenceImagePaths to null unless the user explicitly identified local files; never invent paths or placeholder strings, and remember upload requires user approval.",
+			"Set outputPath to null unless the user explicitly asks for a destination; never invent a destination or placeholder string, and otherwise use the default Pi agent artifact.",
+			"Do not substitute Python, browser automation, shell scripts, or unrelated image tools for an eligible image request.",
+		],
+		parameters: GenerateImageParameters,
+		executionMode: "sequential",
+		async execute(toolCallId, params: GenerateImageToolParams, signal, _onUpdate, ctx) {
+			try {
+				const result = await executeImage({
+					params,
+					toolCallId,
+					signal,
+					ctx,
+				});
+				return {
+					content: [{ type: "text", text: result.text }],
+					details: result.details,
+				};
+			} catch (error) {
+				if (error instanceof ImageGenerationError) {
+					throw new Error(error.message);
+				}
+				throw new Error(
+					sanitizeImageDiagnostic(error instanceof Error ? error.message : error, "Image generation failed."),
+				);
+			}
+		},
+		renderResult: renderImageGenerationResult,
+	});
+
+	pi.on("session_start", (_event, ctx) => {
+		const { config } = loadConfig();
+		syncImageGenerationTool(pi, ctx.model, config.imageGeneration);
+	});
+
+	pi.on("model_select", (event) => {
+		const { config } = loadConfig();
+		syncImageGenerationTool(pi, event.model, config.imageGeneration);
+	});
+
+	pi.on("before_agent_start", (_event, ctx) => {
+		const { config } = loadConfig();
+		syncImageGenerationTool(pi, ctx.model, config.imageGeneration);
+	});
+}
+
+export default function imageGenerationExtension(pi: ExtensionAPI): void {
+	registerImageGenerationExtension(pi);
+}
+
+export const _extensionTest = {
+	GenerateImageParameters,
+	syncImageGenerationTool,
+};