Merge pull request #1064 from unsync/feature/cloudflare-image-generate

feat: adds image generation capability to Workers AI API

Author: VisargD

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,77 @@
+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,
+  },
+  negative_prompt: {
+    param: 'negative_prompt',
+  },
+  steps: [
+    {
+      param: 'num_steps',
+    },
+    {
+      param: 'steps',
+    },
+  ],
+  size: [
+    {
+      param: 'height',
+      transform: (params: any) =>
+        parseInt(params.size.toLowerCase().split('x')[1]),
+    },
+    {
+      param: 'width',
+      transform: (params: any) =>
+        parseInt(params.size.toLowerCase().split('x')[0]),
+    },
+  ],
+  seed: {
+    param: 'seed',
+  },
+};
+
+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;
+};