feat(ai): add xAI image generation (#37778)

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

packages/ai/README.md

@@ -63,6 +63,24 @@ yield *
   })
 ```
 
+xAI image models use the same request API with xAI-native controls:
+
+```ts
+yield *
+  Image.generate({
+    model: XAI.configure({ apiKey }).image("any-model-id"),
+    prompt,
+    options: {
+      n: 2,
+      aspectRatio: "16:9",
+      resolution: "1k",
+      responseFormat: "b64_json",
+      future_option: true,
+    },
+    http,
+  })
+```
+
 Conversational image generation remains part of the LLM interaction. OpenAI Responses exposes it through its hosted image tool:
 
 ```ts

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

@@ -0,0 +1,184 @@
+import { Effect, Encoding, 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,
+  Usage,
+  mergeHttpOptions,
+  mergeJsonRecords,
+  type HttpOptions,
+} from "../schema"
+import { ProviderShared, optionalNull } from "./shared"
+
+const ADAPTER = "xai-images"
+export const DEFAULT_BASE_URL = "https://api.x.ai/v1"
+export const PATH = "/images/generations"
+
+export type XAIImageString<Known extends string> = Known | (string & {})
+
+export type XAIImageOptions = {
+  readonly n?: number
+  readonly aspectRatio?: XAIImageString<
+    | "1:1"
+    | "3:4"
+    | "4:3"
+    | "9:16"
+    | "16:9"
+    | "2:3"
+    | "3:2"
+    | "9:19.5"
+    | "19.5:9"
+    | "9:20"
+    | "20:9"
+    | "1:2"
+    | "2:1"
+    | "auto"
+  >
+  readonly aspect_ratio?: XAIImageString<
+    | "1:1"
+    | "3:4"
+    | "4:3"
+    | "9:16"
+    | "16:9"
+    | "2:3"
+    | "3:2"
+    | "9:19.5"
+    | "19.5:9"
+    | "9:20"
+    | "20:9"
+    | "1:2"
+    | "2:1"
+    | "auto"
+  >
+  readonly resolution?: XAIImageString<"1k" | "2k">
+  readonly responseFormat?: XAIImageString<"url" | "b64_json">
+  readonly response_format?: XAIImageString<"url" | "b64_json">
+} & Record<string, unknown>
+
+type XAIImageBody = Record<string, unknown> & {
+  readonly model: string
+  readonly prompt: string
+}
+
+const XAIImageResponse = Schema.Struct({
+  data: Schema.Array(
+    Schema.Struct({
+      b64_json: optionalNull(Schema.String),
+      url: optionalNull(Schema.String),
+      revised_prompt: optionalNull(Schema.String),
+      mime_type: optionalNull(Schema.String),
+    }),
+  ),
+  usage: Schema.optional(Schema.Unknown),
+})
+
+export interface ModelInput {
+  readonly id: string
+  readonly auth: AuthDefinition
+  readonly baseURL?: string
+  readonly headers?: Record<string, string>
+  readonly http?: HttpOptions
+}
+
+const nativeOptions = (options: XAIImageOptions | undefined) => {
+  if (!options) return undefined
+  const { aspectRatio, responseFormat, ...native } = options
+  return {
+    aspect_ratio: aspectRatio,
+    response_format: responseFormat,
+    ...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<XAIImageOptions> = {
+    id: ADAPTER,
+    generate: Effect.fn("XAIImages.generate")(function* (request: ImageRequestFor<XAIImageOptions>, 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 XAIImageBody
+      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 xAI Images response")),
+      )
+      const decoded = yield* Schema.decodeUnknownEffect(XAIImageResponse)(payload).pipe(
+        Effect.mapError(() => invalidOutput("xAI Images returned an invalid response")),
+      )
+      const images = yield* Effect.forEach(decoded.data, (item, index) => {
+        const mediaType = item.mime_type ?? "application/octet-stream"
+        if (item.b64_json)
+          return Effect.fromResult(Encoding.decodeBase64(item.b64_json)).pipe(
+            Effect.mapError(() => invalidOutput(`xAI Images result ${index} contains invalid base64 data`)),
+            Effect.map(
+              (data) =>
+                new GeneratedImage({
+                  mediaType,
+                  data,
+                  providerMetadata:
+                    item.revised_prompt === undefined || item.revised_prompt === null
+                      ? undefined
+                      : { xai: { revisedPrompt: item.revised_prompt } },
+                }),
+            ),
+          )
+        if (item.url)
+          return Effect.succeed(
+            new GeneratedImage({
+              mediaType,
+              data: item.url,
+              providerMetadata:
+                item.revised_prompt === undefined || item.revised_prompt === null
+                  ? undefined
+                  : { xai: { revisedPrompt: item.revised_prompt } },
+            }),
+          )
+        return Effect.fail(invalidOutput(`xAI Images result ${index} has neither image data nor a URL`))
+      })
+      if (images.length === 0) return yield* invalidOutput("xAI Images returned no images")
+      const usage = ProviderShared.isRecord(decoded.usage) ? decoded.usage : undefined
+      return new ImageResponse({
+        images,
+        usage: usage === undefined ? undefined : new Usage({ providerMetadata: { xai: usage } }),
+        providerMetadata: usage === undefined ? undefined : { xai: { usage } },
+      })
+    }),
+  }
+  return ImageModel.make<XAIImageOptions>({ id: input.id, provider: "xai", route, http: input.http })
+}
+
+export const XAIImages = {
+  model,
+} as const

packages/ai/src/providers/xai.ts

@@ -1,9 +1,10 @@
 import { AuthOptions, type ProviderAuthOption } from "../route/auth-options"
 import type { RouteDefaultsInput } from "../route/client"
-import { ProviderID, type ModelID } from "../schema"
+import { HttpOptions, ProviderID, type ModelID } from "../schema"
 import * as OpenAICompatibleProfiles from "./openai-compatible-profile"
 import * as OpenAICompatibleChat from "../protocols/openai-compatible-chat"
 import * as OpenAIResponses from "../protocols/openai-responses"
+import { XAIImages } from "../protocols/xai-images"
 
 export const id = ProviderID.make("xai")
 
@@ -12,6 +13,8 @@ export type ModelOptions = RouteDefaultsInput &
     readonly baseURL?: string
   }
 
+export type { XAIImageOptions } from "../protocols/xai-images"
+
 export const routes = [OpenAIResponses.route, OpenAICompatibleChat.route]
 
 const auth = (options: ProviderAuthOption<"optional">) => AuthOptions.bearer(options, "XAI_API_KEY")
@@ -41,11 +44,20 @@ export const configure = (input: ModelOptions = {}) => {
   const chatRoute = configuredChatRoute(input)
   const responses = (modelID: string | ModelID) => responsesRoute.model({ id: modelID })
   const chat = (modelID: string | ModelID) => chatRoute.model({ id: modelID })
+  const image = (modelID: string | ModelID) =>
+    XAIImages.model({
+      id: modelID,
+      auth: auth(input),
+      baseURL: input.baseURL ?? OpenAICompatibleProfiles.profiles.xai.baseURL,
+      headers: input.headers,
+      http: input.http === undefined ? undefined : HttpOptions.make(input.http),
+    })
   return {
     id,
     model: responses,
     responses,
     chat,
+    image,
     configure,
   }
 }
@@ -54,3 +66,4 @@ export const provider = configure()
 export const model = provider.model
 export const responses = provider.responses
 export const chat = provider.chat
+export const image = provider.image

packages/ai/test/image.types.ts

@@ -6,7 +6,7 @@ import {
   type ImageRequestFor,
   type ImageRoute,
 } from "../src"
-import { OpenAI } from "../src/providers"
+import { OpenAI, XAI } from "../src/providers"
 
 type GoogleLikeOptions = {
   readonly aspectRatio?: "1:1" | "16:9"
@@ -48,6 +48,30 @@ OpenAI.imageGeneration({ partialImages: "2" })
 // @ts-expect-error Known Google-like options are inferred from the selected model.
 Image.generate({ model: google, prompt: "A lighthouse", options: { aspectRatio: "wide" } })
 
+const xai = XAI.configure({ apiKey: "test" }).image("any-model-id")
+// @ts-expect-error Image generation options are request-scoped, not provider configuration.
+XAI.configure({ image: { options: { resolution: "1k" } } })
+Image.generate({
+  model: xai,
+  prompt: "A lighthouse",
+  options: {
+    n: 2,
+    aspectRatio: "future-ratio",
+    resolution: "future-resolution",
+    responseFormat: "future-format",
+    future_option: true,
+  },
+})
+Image.generate({
+  model: xai,
+  prompt: "A lighthouse",
+  options: { aspect_ratio: "16:9", response_format: "b64_json", native_future_option: true },
+})
+// @ts-expect-error Known xAI numeric options retain their value kind.
+Image.generate({ model: xai, prompt: "A lighthouse", options: { n: "2" } })
+// @ts-expect-error Known xAI string options retain their value kind.
+Image.generate({ model: xai, prompt: "A lighthouse", options: { resolution: 2 } })
+
 declare const generic: ImageModel<ImageOptions>
 Image.generate({ model: generic, prompt: "A lighthouse", options: { arbitrary: true } })
 

packages/ai/test/provider/xai-images.recorded.test.ts

@@ -0,0 +1,33 @@
+import { describe, expect } from "bun:test"
+import { Effect } from "effect"
+import { Image } from "../../src"
+import { XAI } from "../../src/providers"
+import { recordedTests } from "../recorded-test"
+
+const model = XAI.configure({
+  apiKey: process.env.XAI_API_KEY ?? "fixture",
+}).image("grok-imagine-image")
+
+const recorded = recordedTests({
+  prefix: "xai-images",
+  provider: "xai",
+  protocol: "xai-images",
+  requires: ["XAI_API_KEY"],
+})
+
+describe("xAI Images recorded", () => {
+  recorded.effect("generates an image", () =>
+    Effect.gen(function* () {
+      const response = yield* Image.generate({
+        model,
+        prompt: "A simple flat black diamond centered on a plain white background.",
+        options: { aspectRatio: "1:1", resolution: "1k", responseFormat: "b64_json" },
+      })
+
+      expect(response.images).toHaveLength(1)
+      expect(response.image?.mediaType.startsWith("image/")).toBe(true)
+      expect(response.image?.data).toBeInstanceOf(Uint8Array)
+      expect(response.image?.data.length).toBeGreaterThan(0)
+    }),
+  )
+})

packages/ai/test/provider/xai-images.test.ts

@@ -0,0 +1,109 @@
+import { describe, expect } from "bun:test"
+import { Effect, Layer } from "effect"
+import { Headers, HttpClientRequest } from "effect/unstable/http"
+import { Image, ImageClient } from "../../src"
+import { XAI } from "../../src/providers"
+import { Auth } from "../../src/route"
+import { it } from "../lib/effect"
+import { dynamicResponse } from "../lib/http"
+
+describe("xAI Images", () => {
+  it.effect("generates through the OpenAI-compatible Images API", () =>
+    Effect.gen(function* () {
+      const response = yield* Image.generate({
+        model: XAI.configure({
+          apiKey: "test",
+          baseURL: "https://api.xai.test/v1",
+          http: { body: { configured: true }, headers: { "x-default": "yes" } },
+        }).image("grok-imagine-image"),
+        prompt: "A robot tending a rooftop garden",
+        options: {
+          n: 2,
+          aspectRatio: "16:9",
+          aspect_ratio: "4:3",
+          resolution: "1k",
+          responseFormat: "url",
+          response_format: "b64_json",
+          future_option: true,
+        },
+        http: {
+          body: { resolution: "2k", future_option: "http" },
+          headers: { "x-request": "yes" },
+          query: { trace: "1" },
+        },
+      })
+
+      expect(response.images).toHaveLength(2)
+      expect(response.image?.mediaType).toBe("image/jpeg")
+      expect(response.image?.data).toEqual(Uint8Array.from([1, 2, 3]))
+      expect(response.images[1]?.mediaType).toBe("application/octet-stream")
+      expect(response.images[1]?.data).toBe("https://api.xai.test/image.jpg")
+      expect(response.usage?.providerMetadata).toEqual({ xai: { num_images: 2 } })
+      expect(response.providerMetadata).toEqual({ xai: { usage: { num_images: 2 } } })
+    }).pipe(
+      Effect.provide(
+        ImageClient.layer.pipe(
+          Layer.provide(
+            dynamicResponse((input) =>
+              Effect.gen(function* () {
+                const request = yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie)
+                expect(request.url).toBe("https://api.xai.test/v1/images/generations?trace=1")
+                expect(request.headers.get("authorization")).toBe("Bearer test")
+                expect(request.headers.get("x-default")).toBe("yes")
+                expect(request.headers.get("x-request")).toBe("yes")
+                expect(JSON.parse(input.text)).toEqual({
+                  model: "grok-imagine-image",
+                  prompt: "A robot tending a rooftop garden",
+                  n: 2,
+                  aspect_ratio: "4:3",
+                  resolution: "2k",
+                  response_format: "b64_json",
+                  future_option: "http",
+                  configured: true,
+                })
+                return input.respond(
+                  JSON.stringify({
+                    data: [
+                      { b64_json: "AQID", url: null, mime_type: "image/jpeg" },
+                      { b64_json: null, url: "https://api.xai.test/image.jpg", mime_type: null },
+                    ],
+                    usage: { num_images: 2 },
+                  }),
+                  { headers: { "content-type": "application/json" } },
+                )
+              }),
+            ),
+          ),
+        ),
+      ),
+    ),
+  )
+
+  it.effect("supports request-level custom auth", () =>
+    Image.generate({
+      model: XAI.configure({
+        baseURL: "https://api.xai.test/v1",
+        auth: Auth.custom((input) =>
+          Effect.succeed(Headers.set(input.headers, "x-custom-auth", new URL(input.url).hostname)),
+        ),
+      }).image("grok-imagine-image"),
+      prompt: "A robot tending a rooftop garden",
+    }).pipe(
+      Effect.provide(
+        ImageClient.layer.pipe(
+          Layer.provide(
+            dynamicResponse((input) =>
+              Effect.gen(function* () {
+                const request = yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie)
+                expect(request.headers.get("x-custom-auth")).toBe("api.xai.test")
+                return input.respond(JSON.stringify({ data: [{ b64_json: "AQID", mime_type: "image/png" }] }), {
+                  headers: { "content-type": "application/json" },
+                })
+              }),
+            ),
+          ),
+        ),
+      ),
+    ),
+  )
+})