132 lines
4.5 KiB
TypeScript
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
|