feat(ai): add Z.ai image generation (#37780)

Co-authored-by: Aiden Cline <rekram1-node@users.noreply.github.com>

packages/ai/src/protocols/zai-images.ts

@@ -0,0 +1,129 @@
+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, LLMError, mergeHttpOptions, mergeJsonRecords, type HttpOptions } from "../schema"
+import { ProviderShared } from "./shared"
+
+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 LLMError({
+    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) {
+      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

packages/ai/src/providers/index.ts

@@ -15,3 +15,4 @@ export * as OpenAICompatible from "./openai-compatible"
 export * as OpenAICompatibleResponses from "./openai-compatible-responses"
 export * as OpenRouter from "./openrouter"
 export * as XAI from "./xai"
+export * as ZAI from "./zai"

packages/ai/src/providers/zai.ts

@@ -0,0 +1,35 @@
+import { ZAIImages } from "../protocols/zai-images"
+import { AuthOptions, type ProviderAuthOption } from "../route/auth-options"
+import { HttpOptions, ProviderID, type ModelID } from "../schema"
+
+export const id = ProviderID.make("zai")
+
+export type Config = ProviderAuthOption<"optional"> & {
+  readonly baseURL?: string
+  readonly headers?: Record<string, string>
+  readonly http?: HttpOptions.Input
+}
+
+export type { ZAIImageOptions } from "../protocols/zai-images"
+
+const auth = (options: ProviderAuthOption<"optional">) => AuthOptions.bearer(options, "ZAI_API_KEY")
+
+export const configure = (input: Config = {}) => {
+  const image = (modelID: string | ModelID) =>
+    ZAIImages.model({
+      id: modelID,
+      auth: auth(input),
+      baseURL: input.baseURL,
+      headers: input.headers,
+      http: input.http === undefined ? undefined : HttpOptions.make(input.http),
+    })
+
+  return {
+    id,
+    image,
+    configure,
+  }
+}
+
+export const provider = configure()
+export const image = provider.image