opencode/packages/ai/src/protocols/zai-images.ts
2026-08-01 21:57:05 -05:00

132 lines
4.5 KiB
TypeScript

import { Effect, Schema } from "effect"
import { Headers, HttpClientRequest } from "effect/unstable/http"
import { GeneratedImage, ImageModel, ImageResponse, type ImageRequestFor, type ImageRoute } from "../image"
import { Auth, type Definition as AuthDefinition } from "../route/auth"
import { InvalidProviderOutputReason, AIError, mergeHttpOptions, mergeJsonRecords, type HttpOptions } from "../schema"
import { ProviderShared } from "./shared"
import { ImageInputs } from "./utils/image-input"
const ADAPTER = "zai-images"
export const DEFAULT_BASE_URL = "https://api.z.ai/api/paas/v4"
export const PATH = "/images/generations"
export type ZAIImageString<Known extends string> = Known | (string & {})
export type ZAIImageOptions = {
readonly size?: ZAIImageString<
"1024x1024" | "768x1344" | "864x1152" | "1344x768" | "1152x864" | "1440x720" | "720x1440"
>
readonly quality?: ZAIImageString<"hd" | "standard">
readonly userID?: string
} & Record<string, unknown>
type ZAIImageBody = Record<string, unknown> & {
readonly model: string
readonly prompt: string
}
const ZAIImageResponse = Schema.Struct({
created: Schema.optional(Schema.Int),
id: Schema.optional(Schema.String),
request_id: Schema.optional(Schema.String),
data: Schema.Array(Schema.Struct({ url: Schema.String })),
content_filter: Schema.optional(
Schema.Array(
Schema.Struct({
role: Schema.optional(Schema.String),
level: Schema.optional(Schema.Number),
}),
),
),
})
export interface ModelInput {
readonly id: string
readonly auth: AuthDefinition
readonly baseURL?: string
readonly headers?: Record<string, string>
readonly http?: HttpOptions
}
const nativeOptions = (options: ZAIImageOptions | undefined) => {
if (!options) return undefined
const { userID, ...native } = options
return {
user_id: userID,
...native,
}
}
const invalidOutput = (message: string) =>
new AIError({
module: ADAPTER,
method: "generate",
reason: new InvalidProviderOutputReason({ message, route: ADAPTER }),
})
const applyQuery = (url: string, query: Record<string, string> | undefined) => {
if (!query) return url
const next = new URL(url)
Object.entries(query).forEach(([key, value]) => next.searchParams.set(key, value))
return next.toString()
}
export const model = (input: ModelInput) => {
const route: ImageRoute<ZAIImageOptions> = {
id: ADAPTER,
generate: Effect.fn("ZAIImages.generate")(function* (request: ImageRequestFor<ZAIImageOptions>, execute) {
if ((request.images?.length ?? 0) > 0)
return yield* ImageInputs.invalid(ADAPTER, "Z.ai hosted image generation does not support image inputs")
const http = mergeHttpOptions(request.model.http, request.http)
const requestBody = mergeJsonRecords(
{ model: request.model.id, prompt: request.prompt },
nativeOptions(request.options),
http?.body,
) as ZAIImageBody
const text = ProviderShared.encodeJson(requestBody)
const url = applyQuery(`${(input.baseURL ?? DEFAULT_BASE_URL).replace(/\/$/, "")}${PATH}`, http?.query)
const headers = yield* Auth.toEffect(input.auth)({
request,
method: "POST",
url,
body: text,
headers: Headers.fromInput({ ...input.headers, ...http?.headers }),
})
const response = yield* execute(
HttpClientRequest.post(url).pipe(
HttpClientRequest.setHeaders(headers),
HttpClientRequest.bodyText(text, "application/json"),
),
)
const payload = yield* response.json.pipe(
Effect.mapError(() => invalidOutput("Failed to read the Z.ai Images response")),
)
const decoded = yield* Schema.decodeUnknownEffect(ZAIImageResponse)(payload).pipe(
Effect.mapError(() => invalidOutput("Z.ai Images returned an invalid response")),
)
if (decoded.data.length === 0) return yield* invalidOutput("Z.ai Images returned no images")
return new ImageResponse({
images: decoded.data.map(
(item) =>
new GeneratedImage({
mediaType: "application/octet-stream",
data: item.url,
}),
),
providerMetadata: {
zai: {
created: decoded.created,
id: decoded.id,
requestID: decoded.request_id,
contentFilter: decoded.content_filter,
},
},
})
}),
}
return ImageModel.make<ZAIImageOptions>({ id: input.id, provider: "zai", route, http: input.http })
}
export const ZAIImages = {
model,
} as const