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

28.07% Statements 16/57
100% Branches 0/0
0% Functions 0/3
28.07% Lines 16/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 681x           1x 1x   1x 1x 1x 1x 1x               1x 1x 1x 1x 1x           1x 1x   1x                                                                  
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
  AIProviderFunctionConfig,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import type { CreateEmbeddingsRequestBody } from '@shared/types/api/routes/embeddings-api';
import { AIProvider } from '@shared/types/constants';
import { aI21ErrorResponseTransform } from './chat-complete';
 
export const aI21EmbedConfig: AIProviderFunctionConfig = {
  input: {
    param: 'texts',
    required: true,
    transform: (saRequestBody: CreateEmbeddingsRequestBody): string => {
      if ('input' in saRequestBody) {
        return Array.isArray(saRequestBody.input)
          ? saRequestBody.input.join(' ')
          : saRequestBody.input || '';
      }
      throw new Error('Invalid params type for embedding');
    },
  },
  type: {
    param: 'type',
    required: true,
    transform: (saRequestBody: CreateEmbeddingsRequestBody): string => {
      if ('input' in saRequestBody) {
        return 'embed';
      }
      throw new Error('Invalid params type for embedding');
    },
  },
};
 
export const aI21EmbedResponseTransform: ResponseTransformFunction = (
  aiProviderResponseBody,
  aiProviderResponseStatus,
) => {
  if (aiProviderResponseStatus !== 200) {
    const errorResponse = aI21ErrorResponseTransform(aiProviderResponseBody);
    if (errorResponse) return errorResponse;
  }
  if ('results' in aiProviderResponseBody) {
    const results = aiProviderResponseBody.results as {
      embedding: number[];
    }[];
    return {
      object: 'list',
      data: results.map((result, index) => ({
        object: 'embedding' as const,
        embedding: result.embedding,
        index: index,
      })),
      model: '',
      provider: AIProvider.AI21,
      usage: {
        prompt_tokens: -1,
        total_tokens: -1,
      },
    };
  }
 
  return generateInvalidProviderResponseError(
    aiProviderResponseBody,
    AIProvider.AI21,
  );
};