Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions packages/ai/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
184 changes: 184 additions & 0 deletions packages/ai/src/protocols/xai-images.ts
Original file line number Diff line number Diff line change
@@ -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
15 changes: 14 additions & 1 deletion packages/ai/src/providers/xai.ts
Original file line number Diff line number Diff line change
@@ -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")

Expand All @@ -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")
Expand Down Expand Up @@ -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,
}
}
Expand All @@ -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

Large diffs are not rendered by default.

26 changes: 25 additions & 1 deletion packages/ai/test/image.types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 } })

Expand Down
33 changes: 33 additions & 0 deletions packages/ai/test/provider/xai-images.recorded.test.ts
Original file line number Diff line number Diff line change
@@ -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)
}),
)
})
Loading
Loading