All files / api/src/ai-providers/google-vertex-ai transform-generation-config.ts

8.77% Statements 5/57
100% Branches 0/0
0% Functions 0/3
8.77% Lines 5/57

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                  1x                                                                     1x 1x                           1x 1x                                      
import type { EmbedInstancesData } from '@api/ai-providers/google/types';
import type { EmbeddingsParameterTransformFunction } from '@shared/types/api/response/body';
import type { ChatCompletionRequestBody } from '@shared/types/api/routes/chat-completions-api';
import type { CompletionRequestBody } from '@shared/types/api/routes/completions-api';
import type { CreateEmbeddingsRequestBody } from '@shared/types/api/routes/embeddings-api';
 
/**
 * @see https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/gemini#request_body
 */
export const vertexTransformGenerationConfig = (
  params: ChatCompletionRequestBody | CompletionRequestBody,
): Record<
  string,
  string | string[] | number | boolean | Record<string, unknown>
> => {
  const generationConfig: Record<
    string,
    string | string[] | number | boolean | Record<string, unknown>
  > = {};
  if (params.temperature) {
    generationConfig.temperature = params.temperature;
  }
  if (params.top_p) {
    generationConfig.topP = params.top_p;
  }
  // if ('top_k' in params && params.top_k) {
  //   generationConfig.topK = params.top_k;
  // } // TODO: add top_k support
  if (params.max_tokens) {
    generationConfig.maxOutputTokens = params.max_tokens;
  }
  if (params.stop) {
    generationConfig.stopSequences = params.stop;
  }
  if (params.logprobs) {
    generationConfig.responseLogprobs = params.logprobs;
  }
  if (params.top_logprobs) {
    generationConfig.logprobs = params.top_logprobs; // range 1-5, openai supports 1-20
  }
 
  return generationConfig;
};
 
export const googleTransformEmbeddingsDimension: EmbeddingsParameterTransformFunction =
  (params: CreateEmbeddingsRequestBody): Record<string, string | number> => {
    const embeddingsParameters: Record<string, string | number> = {};
    if (params.dimensions) {
      // for multimodal embeddings, the parameter is dimension
      if (Array.isArray(params.input) && typeof params.input[0] === 'object') {
        embeddingsParameters.dimension = params.dimensions;
      } else {
        embeddingsParameters.outputDimensionality = params.dimensions;
      }
    }
 
    return embeddingsParameters;
  };
 
export const googleTransformEmbeddingInput: EmbeddingsParameterTransformFunction =
  (params: CreateEmbeddingsRequestBody): EmbedInstancesData[] => {
    const instances: EmbedInstancesData[] = [];
    if (Array.isArray(params.input)) {
      params.input.forEach((input) => {
        if (typeof input === 'string') {
          instances.push({
            content: input,
            task_type: params.input_type ?? 'text',
          });
        }
      });
    } else {
      instances.push({
        content: params.input,
        task_type: params.input_type ?? 'text',
      });
    }
    return instances;
  };