All files / api/src/ai-providers/google embed.ts

100% Statements 76/76
87.5% Branches 7/8
100% Functions 2/2
100% Lines 76/76

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 85 86 87 88 89 90 91 92 93 94 95 96 971x           1x             1x   1x 1x 1x 1x 1x 2x 2x 2x   2x 2x 1x 2x 2x 2x 1x 1x 1x 1x 1x 1x   2x 2x 2x 2x 1x 1x 1x 1x 1x 1x   1x 1x 1x 1x 1x   1x 3x 3x 3x 3x 3x 3x 3x 1x 1x 1x   2x 2x   3x   3x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x   1x 1x 1x 1x 1x  
import { googleErrorResponseTransform } from '@api/ai-providers/google/chat-complete';
import type {
  GoogleEmbedParams,
  GoogleEmbedResponse,
  GoogleResponseCandidateContent,
} from '@api/ai-providers/google/types';
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
  AIProviderFunctionConfig,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import type { CompletionRequestBody } from '@shared/types/api/routes/completions-api';
import type { CreateEmbeddingsRequestBody } from '@shared/types/api/routes/embeddings-api';
import { AIProvider } from '@shared/types/constants';
 
export const googleEmbedConfig: AIProviderFunctionConfig = {
  input: {
    param: 'content',
    required: true,
    transform: (
      params: CompletionRequestBody | GoogleEmbedParams,
    ): GoogleResponseCandidateContent => {
      const googleParams = params as GoogleEmbedParams;
 
      const parts = [];
      if (Array.isArray(googleParams.input)) {
        googleParams.input.forEach((i) => {
          parts.push({
            text: i,
          });
        });
      } else {
        parts.push({
          text: googleParams.input,
        });
      }
 
      return {
        parts,
      };
    },
  },
  model: {
    param: 'model',
    required: true,
    default: 'embedding-001',
  },
  // Map OpenAI-style 'dimensions' to Google's 'output_dimensionality'
  dimensions: {
    param: 'output_dimensionality',
    required: false,
  },
};
 
export const googleEmbedResponseTransform: ResponseTransformFunction = (
  aiProviderResponseBody,
  aiProviderResponseStatus,
  _responseHeaders,
  _strictOpenAiCompliance,
  saRequestData,
) => {
  if (aiProviderResponseStatus !== 200) {
    const errorResponse = googleErrorResponseTransform(aiProviderResponseBody);
    if (errorResponse) return errorResponse;
  }
 
  const createEmbeddingsRequestBody =
    saRequestData.requestBody as CreateEmbeddingsRequestBody;
 
  const model = (createEmbeddingsRequestBody.model as string) || '';
 
  if ('embedding' in aiProviderResponseBody) {
    const googleResponse =
      aiProviderResponseBody as unknown as GoogleEmbedResponse;
    return {
      object: 'list',
      data: [
        {
          object: 'embedding',
          embedding: googleResponse.embedding.values,
          index: 0,
        },
      ],
      model,
      usage: {
        prompt_tokens: -1,
        total_tokens: -1,
      },
    };
  }
 
  return generateInvalidProviderResponseError(
    aiProviderResponseBody as unknown as Record<string, unknown>,
    AIProvider.GOOGLE,
  );
};