All files / api/src/ai-providers/predibase api.ts

30.43% Statements 7/23
100% Branches 0/0
0% Functions 0/3
30.43% Lines 7/23

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 321x   1x   1x 1x 1x           1x                                   1x  
import { splitString } from '@api/utils/ai-provider';
import type { InternalProviderAPIConfig } from '@shared/types/ai-providers/config';
import { FunctionName } from '@shared/types/api/request';
 
export const predibaseAPIConfig: InternalProviderAPIConfig = {
  getBaseURL: () => 'https://serving.app.predibase.com',
  headers: ({ saTarget: providerOptions }) => {
    return {
      Authorization: `Bearer ${providerOptions.api_key}`,
      Accept: 'application/json',
    };
  },
  getEndpoint: ({ saRequestData }) => {
    // const user = saTarget.user;
    const user = 'predibase'; // TODO: Get from header config
    switch (saRequestData.functionName) {
      case FunctionName.CHAT_COMPLETE: {
        const model = saRequestData.requestBody?.model;
        const base_model = splitString(`${model}`, ':').before;
        /*
        The Predibase model format is "<base_model>[:adapter_id]",
        where adapter_id format is "<adapter_repository_reference/version_number"
        (version_number is required).
        */
        return `/${user}/deployments/v2/llms/${base_model}/v1/chat/completions`;
      }
      default:
        return '';
    }
  },
};