All files / api/src/ai-providers/huggingface complete.ts

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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
import { huggingfaceErrorResponseTransform } from '@api/ai-providers/huggingface/utils';
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
  AIProviderFunctionConfig,
  ResponseChunkStreamTransformFunction,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import type { CompletionResponseBody } from '@shared/types/api/routes/completions-api';
import { AIProvider } from '@shared/types/constants';
import { nanoid } from 'nanoid';
 
// interface HuggingfaceCompleteResponse extends CompletionResponse {}
 
export const HuggingfaceCompleteConfig: AIProviderFunctionConfig = {
  model: {
    param: 'model',
  },
  prompt: {
    param: 'prompt',
    default: '',
  },
  max_tokens: {
    param: 'max_tokens',
    default: 100,
    min: 0,
  },
  temperature: {
    param: 'temperature',
    default: 1,
    min: 0,
    max: 2,
  },
  top_p: {
    param: 'top_p',
    default: 1,
    min: 0,
    max: 1,
  },
  n: {
    param: 'n',
    default: 1,
  },
  stream: {
    param: 'stream',
    default: false,
  },
  logprobs: {
    param: 'logprobs',
    max: 5,
  },
  echo: {
    param: 'echo',
    default: false,
  },
  stop: {
    param: 'stop',
  },
  presence_penalty: {
    param: 'presence_penalty',
    min: -2,
    max: 2,
  },
  frequency_penalty: {
    param: 'frequency_penalty',
    min: -2,
    max: 2,
  },
  best_of: {
    param: 'best_of',
  },
  logit_bias: {
    param: 'logit_bias',
  },
  user: {
    param: 'user',
  },
};
 
export const huggingfaceCompleteResponseTransform: ResponseTransformFunction = (
  aiProviderResponseBody,
  aiProviderResponseStatus,
) => {
  if ('error' in aiProviderResponseBody && aiProviderResponseStatus !== 200) {
    return huggingfaceErrorResponseTransform(
      aiProviderResponseBody,
      aiProviderResponseStatus,
    );
  }
 
  if ('choices' in aiProviderResponseBody) {
    const responseBody = {
      ...aiProviderResponseBody,
    } as unknown as CompletionResponseBody;
    return responseBody;
  }
 
  return generateInvalidProviderResponseError(
    aiProviderResponseBody as unknown as Record<string, unknown>,
    AIProvider.HUGGINGFACE,
  );
};
 
export const huggingfaceCompleteStreamChunkTransform: ResponseChunkStreamTransformFunction =
  (responseChunk) => {
    let chunk = responseChunk.trim();
    if (chunk.startsWith('event: ping')) {
      return '';
    }
 
    chunk = chunk.replace(/^data: /, '');
    chunk = chunk.trim();
    if (chunk === '[DONE]') {
      return 'data: [DONE]\n\n';
    }
    const parsedChunk = JSON.parse(chunk);
    return `data: ${JSON.stringify({
      ...parsedChunk,
      id: `sa-${nanoid()}`,
      provider: AIProvider.HUGGINGFACE,
    })}\n\n`;
  };