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

100% Statements 55/55
91.66% Branches 11/12
100% Functions 1/1
100% Lines 55/55

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 661x         1x   1x   1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x   1x 4x 4x 4x 4x 4x 4x 4x 1x 1x 1x 1x 1x   3x 3x 4x   4x 4x 3x 3x 2x 4x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x   1x 1x 1x 1x 1x  
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
  AIProviderFunctionConfig,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import { AIProvider } from '@shared/types/constants';
import type { MistralAIEmbedResponse } from './types';
import { mistralAIErrorResponseTransform } from './utils';
 
export const mistralAIEmbedConfig: AIProviderFunctionConfig = {
  model: {
    param: 'model',
    required: true,
    default: 'mistral-embed',
  },
  input: {
    param: 'input',
    required: true,
  },
};
 
export const mistralAIEmbedResponseTransform: ResponseTransformFunction = (
  aiProviderResponseBody,
  aiProviderResponseStatus,
  _aiProviderResponseHeaders,
  _strictOpenAiCompliance,
  saRequestData,
) => {
  if (aiProviderResponseStatus !== 200 && 'error' in aiProviderResponseBody) {
    return mistralAIErrorResponseTransform(
      aiProviderResponseBody,
      aiProviderResponseStatus,
    );
  }
 
  const response = aiProviderResponseBody as unknown as MistralAIEmbedResponse;
  const requestBody = saRequestData.requestBody as { model?: string };
  const model = response.model || requestBody.model || '';
 
  if (
    response &&
    typeof response === 'object' &&
    'data' in response &&
    Array.isArray(response.data)
  ) {
    return {
      object: response.object,
      data: response.data.map((d) => ({
        object: d.object,
        embedding: d.embedding,
        index: d.index,
      })),
      model,
      usage: {
        prompt_tokens: response.usage.prompt_tokens,
        total_tokens: response.usage.total_tokens,
      },
    };
  }
 
  return generateInvalidProviderResponseError(
    response as unknown as Record<string, unknown>,
    AIProvider.MISTRAL_AI,
  );
};