feat: adds image generation capability to Workers AI API

Implements support for an image generation feature by adding a new API endpoint.

Includes the configuration and response transformation for handling image generation requests, ensuring proper error handling and response formatting.

Updates the main API configuration to integrate the newly created image generation functionality.

Author: unsync

src/providers/workers-ai/api.ts

@@ -18,8 +18,12 @@ const WorkersAiAPIConfig: ProviderAPIConfig = {
       case 'chatComplete': {
         return `/${model}`;
       }
-      case 'embed':
+      case 'embed': {
         return `/${model}`;
+      }
+      case 'imageGenerate': {
+        return `/${model}`;
+      }
       default:
         return '';
     }

src/providers/workers-ai/chatComplete.ts

@@ -5,10 +5,11 @@ import {
   ErrorResponse,
   ProviderConfig,
 } from '../types';
+import { generateInvalidProviderResponseError } from '../utils';
 import {
-  generateErrorResponse,
-  generateInvalidProviderResponseError,
-} from '../utils';
+  WorkersAiErrorResponse,
+  WorkersAiErrorResponseTransform,
+} from './utils';
 
 export const WorkersAiChatCompleteConfig: ProviderConfig = {
   messages: {
@@ -36,16 +37,6 @@ export const WorkersAiChatCompleteConfig: ProviderConfig = {
   },
 };
 
-export interface WorkersAiErrorObject {
-  code: string;
-  message: string;
-}
-
-interface WorkersAiErrorResponse {
-  success: boolean;
-  errors: WorkersAiErrorObject[];
-}
-
 interface WorkersAiChatCompleteResponse {
   result: {
     response: string;
@@ -60,26 +51,6 @@ interface WorkersAiChatCompleteStreamResponse {
   p?: string;
 }
 
-export const WorkersAiErrorResponseTransform: (
-  response: WorkersAiErrorResponse
-) => ErrorResponse | undefined = (response) => {
-  if ('errors' in response) {
-    return generateErrorResponse(
-      {
-        message: response.errors
-          ?.map((error) => `Error ${error.code}:${error.message}`)
-          .join(', '),
-        type: null,
-        param: null,
-        code: null,
-      },
-      WORKERS_AI
-    );
-  }
-
-  return undefined;
-};
-
 // TODO: cloudflare do not return the usage
 export const WorkersAiChatCompleteResponseTransform: (
   response: WorkersAiChatCompleteResponse | WorkersAiErrorResponse,

src/providers/workers-ai/complete.ts

@@ -1,10 +1,11 @@
 import { Params } from '../../types/requestBody';
 import { CompletionResponse, ErrorResponse, ProviderConfig } from '../types';
 import { WORKERS_AI } from '../../globals';
+import { generateInvalidProviderResponseError } from '../utils';
 import {
-  generateErrorResponse,
-  generateInvalidProviderResponseError,
-} from '../utils';
+  WorkersAiErrorResponse,
+  WorkersAiErrorResponseTransform,
+} from './utils';
 
 export const WorkersAiCompleteConfig: ProviderConfig = {
   prompt: {
@@ -24,16 +25,6 @@ export const WorkersAiCompleteConfig: ProviderConfig = {
   },
 };
 
-export interface WorkersAiErrorObject {
-  code: string;
-  message: string;
-}
-
-interface WorkersAiErrorResponse {
-  success: boolean;
-  errors: WorkersAiErrorObject[];
-}
-
 interface WorkersAiCompleteResponse {
   result: {
     response: string;
@@ -48,26 +39,6 @@ interface WorkersAiCompleteStreamResponse {
   p?: string;
 }
 
-export const WorkersAiErrorResponseTransform: (
-  response: WorkersAiErrorResponse
-) => ErrorResponse | undefined = (response) => {
-  if ('errors' in response) {
-    return generateErrorResponse(
-      {
-        message: response.errors
-          ?.map((error) => `Error ${error.code}:${error.message}`)
-          .join(', '),
-        type: null,
-        param: null,
-        code: null,
-      },
-      WORKERS_AI
-    );
-  }
-
-  return undefined;
-};
-
 export const WorkersAiCompleteResponseTransform: (
   response: WorkersAiCompleteResponse | WorkersAiErrorResponse,
   responseStatus: number,

src/providers/workers-ai/embed.ts

@@ -1,10 +1,11 @@
 import { WORKERS_AI } from '../../globals';
 import { EmbedParams, EmbedResponse } from '../../types/embedRequestBody';
 import { ErrorResponse, ProviderConfig } from '../types';
+import { generateInvalidProviderResponseError } from '../utils';
 import {
-  generateErrorResponse,
-  generateInvalidProviderResponseError,
-} from '../utils';
+  WorkersAiErrorResponse,
+  WorkersAiErrorResponseTransform,
+} from './utils';
 
 export const WorkersAiEmbedConfig: ProviderConfig = {
   input: {
@@ -20,36 +21,6 @@ export const WorkersAiEmbedConfig: ProviderConfig = {
   },
 };
 
-export interface WorkersAiErrorObject {
-  code: string;
-  message: string;
-}
-
-interface WorkersAiErrorResponse {
-  success: boolean;
-  errors: WorkersAiErrorObject[];
-}
-
-export const WorkersAiErrorResponseTransform: (
-  response: WorkersAiErrorResponse
-) => ErrorResponse | undefined = (response) => {
-  if ('errors' in response) {
-    return generateErrorResponse(
-      {
-        message: response.errors
-          ?.map((error) => `Error ${error.code}:${error.message}`)
-          .join(', '),
-        type: null,
-        param: null,
-        code: null,
-      },
-      WORKERS_AI
-    );
-  }
-
-  return undefined;
-};
-
 /**
  * The structure of the CohereEmbedResponse.
  * @interface

src/providers/workers-ai/imageGenerate.ts

@@ -0,0 +1,54 @@
+import { WORKERS_AI } from '../../globals';
+import { ErrorResponse, ImageGenerateResponse, ProviderConfig } from '../types';
+import { generateInvalidProviderResponseError } from '../utils';
+import {
+  WorkersAiErrorResponse,
+  WorkersAiErrorResponseTransform,
+} from './utils';
+
+interface WorkersAiImageGenerateResponse extends ImageGenerateResponse {
+  result: {
+    image?: string;
+  };
+  success: boolean;
+  errors: string[];
+  messages: string[];
+}
+
+export const WorkersAiImageGenerateConfig: ProviderConfig = {
+  prompt: {
+    param: 'prompt',
+    required: true,
+  },
+  steps: {
+    param: 'steps',
+  },
+};
+
+export const WorkersAiImageGenerateResponseTransform: (
+  response: WorkersAiImageGenerateResponse | WorkersAiErrorResponse,
+  responseStatus: number
+) => ImageGenerateResponse | ErrorResponse = (response, responseStatus) => {
+  if (responseStatus !== 200) {
+    const errorResponse = WorkersAiErrorResponseTransform(
+      response as WorkersAiErrorResponse
+    );
+    if (errorResponse) return errorResponse;
+  }
+
+  // the imageGenerate response is not always the same for all Cloudflare image models
+  // we currently only support the image model that returns a base64 image
+  if ('result' in response && 'image' in response.result) {
+    return {
+      created: Math.floor(Date.now() / 1000),
+      data: [
+        {
+          b64_json: response.result.image,
+        },
+      ],
+      provider: WORKERS_AI,
+    };
+  }
+
+  return generateInvalidProviderResponseError(response, WORKERS_AI);
+};

src/providers/workers-ai/index.ts

@@ -11,18 +11,24 @@ import {
   WorkersAiCompleteStreamChunkTransform,
 } from './complete';
 import { WorkersAiEmbedConfig, WorkersAiEmbedResponseTransform } from './embed';
+import {
+  WorkersAiImageGenerateConfig,
+  WorkersAiImageGenerateResponseTransform,
+} from './imageGenerate';
 
 const WorkersAiConfig: ProviderConfigs = {
   complete: WorkersAiCompleteConfig,
   chatComplete: WorkersAiChatCompleteConfig,
   api: WorkersAiAPIConfig,
   embed: WorkersAiEmbedConfig,
+  imageGenerate: WorkersAiImageGenerateConfig,
   responseTransforms: {
     'stream-complete': WorkersAiCompleteStreamChunkTransform,
     complete: WorkersAiCompleteResponseTransform,
     chatComplete: WorkersAiChatCompleteResponseTransform,
     'stream-chatComplete': WorkersAiChatCompleteStreamChunkTransform,
     embed: WorkersAiEmbedResponseTransform,
+    imageGenerate: WorkersAiImageGenerateResponseTransform,
   },
 };
 

src/providers/workers-ai/utils.ts

@@ -0,0 +1,33 @@
+import { ErrorResponse } from '../types';
+import { generateErrorResponse } from '../utils';
+import { WORKERS_AI } from '../../globals';
+
+export interface WorkersAiErrorResponse {
+  success: boolean;
+  errors: WorkersAiErrorObject[];
+}
+
+export interface WorkersAiErrorObject {
+  code: string;
+  message: string;
+}
+
+export const WorkersAiErrorResponseTransform: (
+  response: WorkersAiErrorResponse
+) => ErrorResponse | undefined = (response) => {
+  if ('errors' in response) {
+    return generateErrorResponse(
+      {
+        message: response.errors
+          ?.map((error) => `Error ${error.code}:${error.message}`)
+          .join(', '),
+        type: null,
+        param: null,
+        code: null,
+      },
+      WORKERS_AI
+    );
+  }
+
+  return undefined;
+};