chore: batch & fine-tune improvements

Author: b4s36t4

package-lock.json

@@ -51,9 +51,9 @@
     "async-retry": "^1.3.3",
     "avsc": "^5.7.7",
     "hono": "^4.6.10",
+    "patch-package": "^8.0.0",
     "ws": "^8.18.0",
-    "zod": "^3.22.4",
-    "patch-package": "^8.0.0"
+    "zod": "^3.22.4"
   },
   "devDependencies": {
     "@cloudflare/workers-types": "^4.20230518.0",

package.json

@@ -219,3 +219,9 @@ export const documentMimeTypes = [
   fileExtensionMimeTypeMap.md,
   fileExtensionMimeTypeMap.txt,
 ];
+
+export enum BatchEndpoints {
+  CHAT_COMPLETIONS = '/v1/chat/completions',
+  COMPLETIONS = '/v1/completions',
+  EMBEDDINGS = '/v1/embeddings',
+}

src/globals.ts

@@ -31,6 +31,7 @@ import { ConditionalRouter } from '../services/conditionalRouter';
 import { RouterError } from '../errors/RouterError';
 import { GatewayError } from '../errors/GatewayError';
 import { HookType } from '../middlewares/hooks/types';
+import { Readable } from 'node:stream';
 
 /**
  * Constructs the request options for the API call.
@@ -1058,10 +1059,13 @@ export function constructConfigFromRequestHeaders(
       requestHeaders[`x-${POWERED_BY}-vertex-storage-bucket-name`],
     filename: requestHeaders[`x-${POWERED_BY}-provider-file-name`],
     vertexModelName: requestHeaders[`x-${POWERED_BY}-provider-model`],
+    vertexBatchEndpoint:
+      requestHeaders[`x-${POWERED_BY}-provider-batch-endpoint`],
   };
 
   const fireworksConfig = {
     fireworksAccountId: requestHeaders[`x-${POWERED_BY}-fireworks-account-id`],
+    fireworksFileLength: requestHeaders[`x-${POWERED_BY}-file-upload-size`],
   };
 
   const anthropicConfig = {

src/handlers/handlerUtils.ts

@@ -309,7 +309,7 @@ export function createLineSplitter(): TransformStream {
       leftover = lines.pop() || '';
       for (const line of lines) {
         if (line.trim()) {
-          controller.enqueue(line);
+          controller.enqueue(line.trim());
         }
       }
       return;

src/handlers/streamHandlerUtils.ts

@@ -21,8 +21,23 @@ const FireworksAIAPIConfig: ProviderAPIConfig = {
       Accept: 'application/json',
     };
   },
-  getEndpoint: ({ fn, gatewayRequestBodyJSON: gatewayRequestBody, c }) => {
+  getEndpoint: ({
+    fn,
+    gatewayRequestBodyJSON: gatewayRequestBody,
+    c,
+    gatewayRequestURL,
+  }) => {
     const model = gatewayRequestBody?.model;
+
+    const jobIdIndex = ['cancelFinetune'].includes(fn ?? '') ? -2 : -1;
+    const jobId = gatewayRequestURL.split('/').at(jobIdIndex);
+
+    const url = new URL(gatewayRequestURL);
+    const params = url.searchParams;
+
+    const size = params.get('limit') ?? 50;
+    const page = params.get('after') ?? '1';
+
     switch (fn) {
       case 'complete':
         return '/completions';
@@ -33,7 +48,7 @@ const FireworksAIAPIConfig: ProviderAPIConfig = {
       case 'imageGenerate':
         return `/image_generation/${model}`;
       case 'uploadFile':
-        return `/datasets`;
+        return '';
       case 'retrieveFile': {
         const datasetId = c.req.param('id');
         return `/datasets/${datasetId}`;
@@ -45,13 +60,13 @@ const FireworksAIAPIConfig: ProviderAPIConfig = {
         return `/datasets/${datasetId}`;
       }
       case 'createFinetune':
-        return `/fineTuningJobs`;
+        return `/supervisedFineTuningJobs`;
       case 'retrieveFinetune':
-        return `/fineTuningJobs/${c.req.param('jobId')}`;
+        return `/supervisedFineTuningJobs/${jobId}`;
       case 'listFinetunes':
-        return `/fineTuningJobs`;
+        return `/supervisedFineTuningJobs?pageToken=${page}&pageSize=${size}`;
       case 'cancelFinetune':
-        return `/fineTuningJobs/${c.req.param('jobId')}`;
+        return `/supervisedFineTuningJobs/${jobId}`;
       default:
         return '';
     }

src/providers/fireworks-ai/api.ts

@@ -0,0 +1,84 @@
+import { FIREWORKS_AI } from '../../globals';
+import { Params } from '../../types/requestBody';
+import { RequestHandler } from '../types';
+import FireworksAIAPIConfig from './api';
+import { fireworkFinetuneToOpenAIFinetune } from './utils';
+
+export const FireworkCancelFinetuneResponseTransform = (
+  response: any,
+  status: number
+) => {
+  if (status !== 200) {
+    const error = response?.error || 'Failed to cancel finetune';
+    return new Response(JSON.stringify({ error: { message: error } }), {
+      status: status || 500,
+    });
+  }
+
+  return fireworkFinetuneToOpenAIFinetune(response);
+};
+
+export const FireworksCancelFinetuneRequestHandler: RequestHandler<
+  Params
+> = async ({ requestBody, requestURL, providerOptions, c }) => {
+  const headers = await FireworksAIAPIConfig.headers({
+    c,
+    fn: 'cancelFinetune',
+    providerOptions,
+    transformedRequestUrl: requestURL,
+    transformedRequestBody: requestBody,
+  });
+
+  const baseURL = await FireworksAIAPIConfig.getBaseURL({
+    c,
+    gatewayRequestURL: requestURL,
+    providerOptions,
+  });
+
+  const endpoint = FireworksAIAPIConfig.getEndpoint({
+    c,
+    fn: 'cancelFinetune',
+    gatewayRequestBodyJSON: requestBody,
+    gatewayRequestURL: requestURL,
+    providerOptions,
+  });
+
+  try {
+    const request = await fetch(baseURL + endpoint, {
+      method: 'DELETE',
+      headers,
+      body: JSON.stringify(requestBody),
+    });
+
+    if (!request.ok) {
+      const error = await request.json();
+      return new Response(
+        JSON.stringify({
+          error: { message: (error as any).error },
+          provider: FIREWORKS_AI,
+        }),
+        {
+          status: 500,
+          headers: {
+            'Content-Type': 'application/json',
+          },
+        }
+      );
+    }
+
+    const response = await request.json();
+
+    const mappedResponse = fireworkFinetuneToOpenAIFinetune(response as any);
+
+    return new Response(JSON.stringify(mappedResponse), {
+      status: 200,
+      headers: { 'Content-Type': 'application/json' },
+    });
+  } catch (error) {
+    const errorMessage =
+      error instanceof Error ? error.message : 'Unknown error';
+    return new Response(JSON.stringify({ error: { message: errorMessage } }), {
+      status: 500,
+    });
+  }
+};