All files / api/src/ai-providers/google-vertex-ai create-batch.ts

36.76% Statements 25/68
100% Branches 0/0
0% Functions 0/6
36.76% Lines 25/68

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84                        1x   1x 1x 1x 1x 1x               1x 1x 1x 1x 1x               1x 1x 1x 1x 1x                       1x                   1x 1x 1x 1x 1x     1x 1x   1x                  
import type {
  GoogleBatchRecord,
  GoogleBatchRecordOutputConfig,
} from '@api/ai-providers/google/types';
 
import type {
  AIProviderFunctionConfig,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
 
import type { CreateBatchResponseBody } from '@shared/types/api/routes/batch-api';
import type { CreateEmbeddingsRequestBody } from '@shared/types/api/routes/embeddings-api';
import { GoogleToOpenAIBatch, getModelAndProvider } from './utils';
 
export const GoogleBatchCreateConfig: AIProviderFunctionConfig = {
  model: {
    param: 'model',
    required: true,
    transform: (saRequestBody: CreateEmbeddingsRequestBody): string => {
      if (!saRequestBody.model) {
        throw new Error('Model is required');
      }
 
      const { model, provider } = getModelAndProvider(saRequestBody.model);
      return `publishers/${provider}/models/${model}`;
    },
  },
  input_file_id: {
    param: 'inputConfig',
    required: true,
    transform: (params: Record<string, unknown>) => {
      return {
        instancesFormat: 'jsonl',
        gcsSource: {
          uris: decodeURIComponent(params.input_file_id as string),
        },
      };
    },
  },
  output_data_config: {
    param: 'outputConfig',
    required: true,
    transform: (
      params: Record<string, unknown>,
    ): GoogleBatchRecordOutputConfig => {
      return {
        predictionsFormat: 'jsonl',
        gcsDestination: {
          outputUriPrefix: decodeURIComponent(
            params.output_data_config as string,
          ),
        },
      };
    },
    default: (params: Record<string, unknown>) => {
      const inputFileId = decodeURIComponent(params.input_file_id as string);
      const gcsURLToContainingFolder = `${inputFileId.split('/').slice(0, -1).join('/')}/`;
      return {
        predictionsFormat: 'jsonl',
        gcsDestination: {
          outputUriPrefix: gcsURLToContainingFolder,
        },
      };
    },
  },
  job_name: {
    param: 'displayName',
    required: true,
    default: () => {
      return crypto.randomUUID();
    },
  },
};
 
export const googleBatchCreateResponseTransform: ResponseTransformFunction = (
  response,
  status,
) => {
  if (status === 200) {
    return GoogleToOpenAIBatch(response as unknown as GoogleBatchRecord);
  }
  return response as unknown as CreateBatchResponseBody;
};