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

100% Statements 58/58
90.9% Branches 10/11
100% Functions 1/1
100% Lines 58/58

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 691x         1x   1x   1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x   1x 5x 5x 5x 5x 5x 5x 5x 1x 1x 1x 1x 1x   4x 4x 5x   5x 5x 4x 4x 3x 5x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x   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 { OllamaEmbedResponse } from './types';
import { ollamaErrorResponseTransform } from './utils';
 
export const ollamaEmbedConfig: AIProviderFunctionConfig = {
  model: {
    param: 'model',
    required: true,
    default: 'llama3.2:latest',
  },
  input: {
    param: 'prompt',
    required: true,
  },
};
 
export const ollamaEmbedResponseTransform: ResponseTransformFunction = (
  aiProviderResponseBody,
  aiProviderResponseStatus,
  _aiProviderResponseHeaders,
  _strictOpenAiCompliance,
  saRequestData,
) => {
  if (aiProviderResponseStatus !== 200 && 'error' in aiProviderResponseBody) {
    return ollamaErrorResponseTransform(
      aiProviderResponseBody,
      aiProviderResponseStatus,
    );
  }
 
  const response = aiProviderResponseBody as unknown as OllamaEmbedResponse;
  const requestBody = saRequestData.requestBody as { model?: string };
  const model = response.model || requestBody.model || '';
 
  if (
    response &&
    typeof response === 'object' &&
    'embedding' in response &&
    Array.isArray(response.embedding)
  ) {
    return {
      object: 'list',
      data: [
        {
          object: 'embedding',
          embedding: response.embedding,
          index: 0,
        },
      ],
      model,
      usage: {
        prompt_tokens: response.prompt_eval_count ?? 0,
        total_tokens:
          (response.prompt_eval_count ?? 0) + (response.eval_count ?? 0),
      },
    };
  }
 
  return generateInvalidProviderResponseError(
    response as unknown as Record<string, unknown>,
    AIProvider.OLLAMA,
  );
};