chore: vertex finetune and batches cleanup

src/handlers/handlerUtils.ts

@@ -18,11 +18,7 @@ import {
   FIREWORKS_AI,
 } from '../globals';
 import Providers from '../providers';
-import {
-  ProviderAPIConfig,
-  RequestHandlers,
-  endpointStrings,
-} from '../providers/types';
+import { ProviderAPIConfig, endpointStrings } from '../providers/types';
 import transformToProviderRequest from '../services/transformToProviderRequest';
 import { Options, Params, StrategyModes, Targets } from '../types/requestBody';
 import { convertKeysToCamelCase } from '../utils';

src/providers/bedrock/cancelBatch.ts

@@ -10,10 +10,10 @@ export const BedrockCancelBatchResponseTransform = (
   gatewayRequestUrl: string
 ): CancelBatchResponse | ErrorResponse => {
   if (responseStatus !== 200) {
-    const errorResposne = BedrockErrorResponseTransform(
+    const errorResponse = BedrockErrorResponseTransform(
       response as BedrockErrorResponse
     );
-    if (errorResposne) return errorResposne;
+    if (errorResponse) return errorResponse;
   }
   const batchId = decodeURIComponent(
     gatewayRequestUrl.split('/v1/batches/')[1].split('/')[0]

src/providers/bedrock/createBatch.ts

@@ -65,10 +65,10 @@ export const BedrockCreateBatchResponseTransform: (
   responseStatus: number
 ) => CreateBatchResponse | ErrorResponse = (response, responseStatus) => {
   if (responseStatus !== 200) {
-    const errorResposne = BedrockErrorResponseTransform(
+    const errorResponse = BedrockErrorResponseTransform(
       response as BedrockErrorResponse
     );
-    if (errorResposne) return errorResposne;
+    if (errorResponse) return errorResponse;
   }
 
   if ('jobArn' in response) {

src/providers/bedrock/createFinetune.ts

@@ -1,6 +1,7 @@
 import { ErrorResponse, FinetuneRequest, ProviderConfig } from '../types';
 import { BedrockErrorResponseTransform } from './chatComplete';
 import { BedrockErrorResponse } from './embed';
+import { populateHyperParameters } from './utils';
 
 export const BedrockCreateFinetuneConfig: ProviderConfig = {
   model: {
@@ -15,10 +16,10 @@ export const BedrockCreateFinetuneConfig: ProviderConfig = {
     param: 'hyperParameters',
     required: false,
     transform: (value: FinetuneRequest) => {
-      const epochCount = value.hyperparameters?.n_epochs;
-      const learningRateMultiplier =
-        value.hyperparameters?.learning_rate_multiplier;
-      const batchSize = value.hyperparameters?.batch_size;
+      const hyperParameters = populateHyperParameters(value);
+      const epochCount = hyperParameters.n_epochs;
+      const learningRateMultiplier = hyperParameters.learning_rate_multiplier;
+      const batchSize = hyperParameters.batch_size;
       return {
         epochCount: epochCount ? String(epochCount) : undefined,
         learningRateMultiplier: learningRateMultiplier

src/providers/bedrock/retrieveBatch.ts

@@ -11,10 +11,10 @@ export const BedrockRetrieveBatchResponseTransform = (
   responseStatus: number
 ): RetrieveBatchResponse | ErrorResponse => {
   if (responseStatus !== 200) {
-    const errorResposne = BedrockErrorResponseTransform(
+    const errorResponse = BedrockErrorResponseTransform(
       response as BedrockErrorResponse
     );
-    if (errorResposne) return errorResposne;
+    if (errorResponse) return errorResponse;
   }
 
   if ('jobArn' in response) {

src/providers/bedrock/uploadFile.ts

@@ -383,7 +383,7 @@ export const BedrockUploadFileRequestHandler: RequestHandler<
       filename: s3Url,
       purpose,
       bytes: handler.contentLength,
-      status: 'uploaded',
+      status: 'processed',
       status_details: '',
     };
 

src/providers/bedrock/uploadFileUtils.ts

@@ -57,7 +57,7 @@ interface AnthropicMessage extends Message {
   content?: string | AnthropicMessageContentItem[];
 }
 
-interface AnthorpicTextContentItem {
+interface AnthropicTextContentItem {
   type: 'text';
   text: string;
 }
@@ -69,7 +69,7 @@ interface AnthropicToolContentItem {
   input: Record<string, any>;
 }
 
-type AnthropicContentItem = AnthorpicTextContentItem | AnthropicToolContentItem;
+type AnthropicContentItem = AnthropicTextContentItem | AnthropicToolContentItem;
 
 const transformAssistantMessageForAnthropic = (
   msg: Message

src/providers/bedrock/utils.ts

@@ -11,6 +11,7 @@ import {
 import { Options } from '../../types/requestBody';
 import { GatewayError } from '../../errors/GatewayError';
 import { BedrockFinetuneRecord } from './types';
+import { FinetuneRequest } from '../types';
 
 export const generateAWSHeaders = async (
   body: Record<string, any> | string | undefined,
@@ -369,3 +370,14 @@ export async function providerAssumedRoleCredentials(
     throw new GatewayError('Error while assuming bedrock role');
   }
 }
+
+export const populateHyperParameters = (value: FinetuneRequest) => {
+  let hyperParameters = value.hyperparameters ?? {};
+
+  if (value.method) {
+    const method = value.method.type;
+    hyperParameters = value.method?.[method]?.hyperparameters ?? {};
+  }
+
+  return hyperParameters;
+};

src/providers/google-vertex-ai/api.ts

@@ -1,6 +1,6 @@
 import { Options } from '../../types/requestBody';
 import { endpointStrings, ProviderAPIConfig } from '../types';
-import { getModelAndProvider, getAccessToken } from './utils';
+import { getModelAndProvider, getAccessToken, getBucketAndFile } from './utils';
 
 const getApiVersion = (provider: string, inputModel: string) => {
   if (provider === 'meta') return 'v1beta1';
@@ -40,6 +40,10 @@ const BATCH_ENDPOINTS = [
   'getBatchOutput',
   'listBatches',
   'cancelBatch',
+  'createFinetune',
+  'retrieveFinetune',
+  'listFinetunes',
+  'cancelFinetune',
 ];
 const NON_INFERENCE_ENDPOINTS = [...FILE_ENDPOINTS, ...BATCH_ENDPOINTS];
 
@@ -82,7 +86,13 @@ export const GoogleApiConfig: ProviderAPIConfig = {
     }
 
     if (NON_INFERENCE_ENDPOINTS.includes(fn)) {
-      const jobIdIndex = ['cancelBatch'].includes(fn) ? -2 : -1;
+      const jobIdIndex = [
+        'cancelBatch',
+        'retrieveFileContent',
+        'cancelFinetune',
+      ].includes(fn)
+        ? -2
+        : -1;
       const jobId = gatewayRequestURL.split('/').at(jobIdIndex);
 
       const url = new URL(gatewayRequestURL);
@@ -104,14 +114,29 @@ export const GoogleApiConfig: ProviderAPIConfig = {
           return `/v1/projects/${projectId}/locations/${vertexRegion}/batchPredictionJobs/${jobId}:cancel`;
         }
         case 'uploadFile':
+        case 'getBatchOutput':
           // We handle file upload in a separate request handler
           return '';
+        case 'retrieveFile':
+          return '';
+        case 'retrieveFileContent': {
+          const { bucket, file } = getBucketAndFile(jobId ?? '');
+          return `/${bucket}/${file}`;
+        }
         case 'createBatch':
           return `/v1/projects/${projectId}/locations/${vertexRegion}/batchPredictionJobs`;
-        case 'getBatchOutput':
-          return '';
-        default:
-          return '';
+        case 'createFinetune':
+          return `/v1/projects/${projectId}/locations/${vertexRegion}/tuningJobs`;
+        case 'listFinetunes': {
+          const pageSize = searchParams.get('limit') ?? 20;
+          const after = searchParams.get('after') ?? '';
+          return `/v1/projects/${projectId}/locations/${vertexRegion}/tuningJobs?pageSize=${pageSize}&pageToken=${after}`;
+        }
+        case 'retrieveFinetune':
+          return `/v1/projects/${projectId}/locations/${vertexRegion}/tuningJobs/${jobId}`;
+        case 'cancelFinetune': {
+          return `/v1/projects/${projectId}/locations/${vertexRegion}/tuningJobs/${jobId}:cancel`;
+        }
       }
     }
 

src/providers/google-vertex-ai/createFinetune.ts

@@ -0,0 +1,45 @@
+import { GOOGLE_VERTEX_AI } from '../../globals';
+import { ProviderConfig } from '../types';
+import { GoogleErrorResponse, GoogleFinetuneRecord } from './types';
+import { GoogleToOpenAIFinetune, transformVertexFinetune } from './utils';
+
+export const GoogleVertexFinetuneConfig: ProviderConfig = {
+  model: {
+    param: 'baseModel',
+    required: true,
+  },
+  training_file: {
+    param: 'supervisedTuningSpec',
+    required: true,
+    transform: transformVertexFinetune,
+  },
+  suffix: {
+    param: 'tunedModelDisplayName',
+    required: true,
+  },
+  validation_file: {
+    param: 'supervisedTuningSpec',
+    required: false,
+    transform: transformVertexFinetune,
+  },
+  method: {
+    param: 'supervisedTuningSpec',
+    required: false,
+    transform: transformVertexFinetune,
+  },
+  hyperparameters: {
+    param: 'supervisedTuningSpec',
+    required: false,
+    transform: transformVertexFinetune,
+  },
+};
+
+export const GoogleFinetuneCreateResponseTransform = (
+  input: Response | GoogleErrorResponse,
+  status: number
+) => {
+  if (status !== 200) {
+    return { ...input, provider: GOOGLE_VERTEX_AI };
+  }
+  return GoogleToOpenAIFinetune(input as unknown as GoogleFinetuneRecord);
+};

src/providers/google-vertex-ai/index.ts

@@ -35,6 +35,18 @@ import {
   GoogleFileUploadResponseTransform,
 } from './uploadFile';
 import { GoogleRetrieveBatchResponseTransform } from './retrieveBatch';
+import {
+  GoogleFinetuneCreateResponseTransform,
+  GoogleVertexFinetuneConfig,
+} from './createFinetune';
+import { GoogleRetrieveFileContentResponseTransform } from './retrieveFileContent';
+import {
+  GoogleRetrieveFileRequestHandler,
+  GoogleRetrieveFileResponseTransform,
+} from './retrieveFile';
+import { GoogleFinetuneRetrieveResponseTransform } from './retrieveFinetune';
+import { GoogleFinetuneListResponseTransform } from './listFinetunes';
+import { GoogleListFilesRequestHandler } from './listFiles';
 
 const VertexConfig: ProviderConfigs = {
   api: VertexApiConfig,
@@ -45,6 +57,10 @@ const VertexConfig: ProviderConfigs = {
       retrieveBatch: {},
       listBatches: {},
       cancelBatch: {},
+      createFinetune: GoogleVertexFinetuneConfig,
+      retrieveFile: {},
+      cancelFinetune: {},
+      retrieveFileContent: {},
     };
 
     const responseTransforms = {
@@ -54,6 +70,11 @@ const VertexConfig: ProviderConfigs = {
       listBatches: GoogleListBatchesResponseTransform,
       cancelBatch: GoogleCancelBatchResponseTransform,
       createBatch: GoogleBatchCreateResponseTransform,
+      retrieveFileContent: GoogleRetrieveFileContentResponseTransform,
+      retrieveFile: GoogleRetrieveFileResponseTransform,
+      createFinetune: GoogleFinetuneCreateResponseTransform,
+      retrieveFinetune: GoogleFinetuneRetrieveResponseTransform,
+      listFinetunes: GoogleFinetuneListResponseTransform,
     };
 
     const baseConfig = {
@@ -76,6 +97,7 @@ const VertexConfig: ProviderConfigs = {
           embed: GoogleEmbedConfig,
           imageGenerate: GoogleImageGenConfig,
           createBatch: GoogleBatchCreateConfig,
+          createFinetune: baseConfig.createFinetune,
           responseTransforms: {
             'stream-chatComplete': GoogleChatCompleteStreamChunkTransform,
             chatComplete: GoogleChatCompleteResponseTransform,
@@ -89,6 +111,7 @@ const VertexConfig: ProviderConfigs = {
           chatComplete: VertexAnthropicChatCompleteConfig,
           api: GoogleApiConfig,
           createBatch: GoogleBatchCreateConfig,
+          createFinetune: baseConfig.createFinetune,
           responseTransforms: {
             'stream-chatComplete':
               VertexAnthropicChatCompleteStreamChunkTransform,
@@ -101,6 +124,7 @@ const VertexConfig: ProviderConfigs = {
           chatComplete: VertexLlamaChatCompleteConfig,
           createBatch: GoogleBatchCreateConfig,
           api: GoogleApiConfig,
+          createFinetune: baseConfig.createFinetune,
           responseTransforms: {
             chatComplete: VertexLlamaChatCompleteResponseTransform,
             'stream-chatComplete': VertexLlamaChatCompleteStreamChunkTransform,
@@ -114,6 +138,7 @@ const VertexConfig: ProviderConfigs = {
           }),
           createBatch: GoogleBatchCreateConfig,
           api: GoogleApiConfig,
+          createFinetune: baseConfig.createFinetune,
           responseTransforms: {
             ...responseTransformers(GOOGLE_VERTEX_AI, {
               chatComplete: true,
@@ -128,6 +153,8 @@ const VertexConfig: ProviderConfigs = {
   requestHandlers: {
     uploadFile: GoogleFileUploadRequestHandler,
     getBatchOutput: BatchOutputRequestHandler,
+    listFiles: GoogleListFilesRequestHandler,
+    retrieveFile: GoogleRetrieveFileRequestHandler,
   },
 };
 

src/providers/google-vertex-ai/listFiles.ts

@@ -0,0 +1,12 @@
+import { GOOGLE_VERTEX_AI } from '../../globals';
+
+export const GoogleListFilesRequestHandler = async () => {
+  return new Response(
+    JSON.stringify({
+      message: 'listFiles is not supported by Google Vertex AI',
+      status: 'failure',
+      provider: GOOGLE_VERTEX_AI,
+    }),
+    { status: 500, headers: { 'Content-Type': 'application/json' } }
+  );
+};

src/providers/google-vertex-ai/listFinetunes.ts

@@ -0,0 +1,31 @@
+import { GOOGLE_VERTEX_AI } from '../../globals';
+import { GoogleFinetuneRecord, GoogleErrorResponse } from './types';
+import { GoogleToOpenAIFinetune } from './utils';
+
+type GoogleListFinetunesResponse = {
+  tuningJobs: GoogleFinetuneRecord[];
+  nextPageToken?: string;
+};
+
+export const GoogleFinetuneListResponseTransform = (
+  input: Response | GoogleErrorResponse,
+  status: number
+) => {
+  if (status !== 200) {
+    return { ...input, provider: GOOGLE_VERTEX_AI };
+  }
+  const records =
+    (input as unknown as { tuningJobs: GoogleFinetuneRecord[] }).tuningJobs ??
+    [];
+
+  const objects = records.map(GoogleToOpenAIFinetune);
+
+  return {
+    data: objects,
+    object: 'list',
+    first_id: objects.at(0)?.id,
+    last_id: objects.at(-1)?.id,
+    has_more: !!(input as unknown as GoogleListFinetunesResponse)
+      ?.nextPageToken,
+  };
+};