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

42.85% Statements 18/42
100% Branches 0/0
0% Functions 0/1
42.85% Lines 18/42

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  1x         1x   1x   1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x   1x                                                      
import type { AnyscaleEmbedResponse } from '@api/ai-providers/anyscale/types';
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
  AIProviderFunctionConfig,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import { AIProvider } from '@shared/types/constants';
 
import { anyscaleErrorResponseTransform } from './chat-complete';
 
export const anyscaleEmbedConfig: AIProviderFunctionConfig = {
  model: {
    param: 'model',
    required: true,
    default: 'thenlper/gte-large',
  },
  input: {
    param: 'input',
    default: '',
  },
  user: {
    param: 'user',
  },
};
 
export const anyscaleEmbedResponseTransform: ResponseTransformFunction = (
  aiProviderResponseBody,
  aiProviderResponseStatus,
) => {
  if (aiProviderResponseStatus !== 200) {
    const errorResponse = anyscaleErrorResponseTransform(
      aiProviderResponseBody,
    );
    if (errorResponse) return errorResponse;
  }
 
  if ('data' in aiProviderResponseBody) {
    const createEmbeddingsResponseBody =
      aiProviderResponseBody as AnyscaleEmbedResponse;
    return {
      object: createEmbeddingsResponseBody.object,
      data: createEmbeddingsResponseBody.data,
      model: createEmbeddingsResponseBody.model,
      usage: createEmbeddingsResponseBody.usage,
    };
  }
 
  return generateInvalidProviderResponseError(
    aiProviderResponseBody,
    AIProvider.ANYSCALE,
  );
};