-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
friendli.ts
371 lines (330 loc) Β· 10.3 KB
/
friendli.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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
import { CallbackManagerForLLMRun } from "@langchain/core/callbacks/manager";
import {
BaseChatModel,
BaseChatModelCallOptions,
type BaseChatModelParams,
} from "@langchain/core/language_models/chat_models";
import {
BaseMessage,
AIMessage,
ChatMessage,
HumanMessage,
SystemMessage,
HumanMessageChunk,
AIMessageChunk,
SystemMessageChunk,
ChatMessageChunk,
} from "@langchain/core/messages";
import {
ChatGeneration,
ChatGenerationChunk,
ChatResult,
} from "@langchain/core/outputs";
import { getEnvironmentVariable } from "@langchain/core/utils/env";
import { convertEventStreamToIterableReadableDataStream } from "../utils/event_source_parse.js";
/**
* Type representing the role of a message in the Friendli chat model.
*/
export type FriendliMessageRole = "system" | "assistant" | "user";
interface FriendliMessage {
role: FriendliMessageRole;
content: string;
}
function messageToFriendliRole(message: BaseMessage): FriendliMessageRole {
const type = message._getType();
switch (type) {
case "ai":
return "assistant";
case "human":
return "user";
case "system":
return "system";
case "function":
throw new Error("Function messages not supported");
case "generic": {
if (!ChatMessage.isInstance(message)) {
throw new Error("Invalid generic chat message");
}
if (["system", "assistant", "user"].includes(message.role)) {
return message.role as FriendliMessageRole;
}
throw new Error(`Unknown message type: ${type}`);
}
default:
throw new Error(`Unknown message type: ${type}`);
}
}
function friendliResponseToChatMessage(message: FriendliMessage): BaseMessage {
switch (message.role) {
case "user":
return new HumanMessage(message.content ?? "");
case "assistant":
return new AIMessage(message.content ?? "");
case "system":
return new SystemMessage(message.content ?? "");
default:
return new ChatMessage(message.content ?? "", message.role ?? "unknown");
}
}
function _convertDeltaToMessageChunk(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
delta: Record<string, any>
) {
const role = delta.role ?? "assistant";
const content = delta.content ?? "";
let additional_kwargs;
if (delta.function_call) {
additional_kwargs = {
function_call: delta.function_call,
};
} else {
additional_kwargs = {};
}
if (role === "user") {
return new HumanMessageChunk({ content });
} else if (role === "assistant") {
return new AIMessageChunk({ content, additional_kwargs });
} else if (role === "system") {
return new SystemMessageChunk({ content });
} else {
return new ChatMessageChunk({ content, role });
}
}
/**
* The ChatFriendliParams interface defines the input parameters for
* the ChatFriendli class.
*/
export interface ChatFriendliParams extends BaseChatModelParams {
/**
* Model name to use.
*/
model?: string;
/**
* Base endpoint url.
*/
baseUrl?: string;
/**
* Friendli personal access token to run as.
*/
friendliToken?: string;
/**
* Friendli team ID to run as.
*/
friendliTeam?: string;
/**
* Number between -2.0 and 2.0. Positive values penalizes tokens that have been
* sampled, taking into account their frequency in the preceding text. This
* penalization diminishes the model's tendency to reproduce identical lines
* verbatim.
*/
frequencyPenalty?: number;
/**
* Number between -2.0 and 2.0. Positive values penalizes tokens that have been
* sampled at least once in the existing text.
* presence_penalty: Optional[float] = None
* The maximum number of tokens to generate. The length of your input tokens plus
* `max_tokens` should not exceed the model's maximum length (e.g., 2048 for OpenAI
* GPT-3)
*/
maxTokens?: number;
/**
* When one of the stop phrases appears in the generation result, the API will stop
* generation. The phrase is included in the generated result. If you are using
* beam search, all of the active beams should contain the stop phrase to terminate
* generation. Before checking whether a stop phrase is included in the result, the
* phrase is converted into tokens.
*/
stop?: string[];
/**
* Sampling temperature. Smaller temperature makes the generation result closer to
* greedy, argmax (i.e., `top_k = 1`) sampling. If it is `None`, then 1.0 is used.
*/
temperature?: number;
/**
* Tokens comprising the top `top_p` probability mass are kept for sampling. Numbers
* between 0.0 (exclusive) and 1.0 (inclusive) are allowed. If it is `None`, then 1.0
* is used by default.
*/
topP?: number;
/**
* Additional kwargs to pass to the model.
*/
modelKwargs?: Record<string, unknown>;
}
/**
* The ChatFriendli class is used to interact with Friendli inference Endpoint models.
* This requires your Friendli Token and Friendli Team which is autoloaded if not specified.
*/
export class ChatFriendli extends BaseChatModel<BaseChatModelCallOptions> {
lc_serializable = true;
static lc_name() {
return "Friendli";
}
get lc_secrets(): { [key: string]: string } | undefined {
return {
friendliToken: "FRIENDLI_TOKEN",
friendliTeam: "FRIENDLI_TEAM",
};
}
model = "llama-2-13b-chat";
baseUrl = "https://inference.friendli.ai";
friendliToken?: string;
friendliTeam?: string;
frequencyPenalty?: number;
maxTokens?: number;
stop?: string[];
temperature?: number;
topP?: number;
modelKwargs?: Record<string, unknown>;
constructor(fields: ChatFriendliParams) {
super(fields);
this.model = fields?.model ?? this.model;
this.baseUrl = fields?.baseUrl ?? this.baseUrl;
this.friendliToken =
fields?.friendliToken ?? getEnvironmentVariable("FRIENDLI_TOKEN");
this.friendliTeam =
fields?.friendliTeam ?? getEnvironmentVariable("FRIENDLI_TEAM");
this.frequencyPenalty = fields?.frequencyPenalty ?? this.frequencyPenalty;
this.maxTokens = fields?.maxTokens ?? this.maxTokens;
this.stop = fields?.stop ?? this.stop;
this.temperature = fields?.temperature ?? this.temperature;
this.topP = fields?.topP ?? this.topP;
this.modelKwargs = fields?.modelKwargs ?? {};
if (!this.friendliToken) {
throw new Error("Missing Friendli Token");
}
}
_llmType() {
return "friendli";
}
private constructHeaders(stream: boolean) {
return {
"Content-Type": "application/json",
Accept: stream ? "text/event-stream" : "application/json",
Authorization: `Bearer ${this.friendliToken}`,
"X-Friendli-Team": this.friendliTeam ?? "",
};
}
private constructBody(
messages: BaseMessage[],
stream: boolean,
_options?: this["ParsedCallOptions"]
) {
const messageList = messages.map((message) => {
if (typeof message.content !== "string") {
throw new Error(
"Friendli does not support non-string message content."
);
}
return {
role: messageToFriendliRole(message),
content: message.content,
};
});
const body = JSON.stringify({
messages: messageList,
stream,
model: this.model,
max_tokens: this.maxTokens,
frequency_penalty: this.frequencyPenalty,
stop: this.stop,
temperature: this.temperature,
top_p: this.topP,
...this.modelKwargs,
});
return body;
}
/**
* Calls the Friendli endpoint and retrieves the result.
* @param {BaseMessage[]} messages The input messages.
* @returns {Promise<ChatResult>} A promise that resolves to the generated chat result.
*/
/** @ignore */
async _generate(
messages: BaseMessage[],
_options: this["ParsedCallOptions"]
): Promise<ChatResult> {
interface ChatFriendliResponse {
choices: {
index: number;
message: {
role: FriendliMessageRole;
content: string;
};
finish_reason: string;
}[];
usage: {
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;
};
created: number;
}
const response = (await this.caller.call(async () =>
fetch(`${this.baseUrl}/v1/chat/completions`, {
method: "POST",
headers: this.constructHeaders(false),
body: this.constructBody(messages, false, _options),
}).then((res) => res.json())
)) as ChatFriendliResponse;
const generations: ChatGeneration[] = [];
for (const data of response.choices ?? []) {
const text = data.message?.content ?? "";
const generation: ChatGeneration = {
text,
message: friendliResponseToChatMessage(data.message ?? {}),
};
if (data.finish_reason) {
generation.generationInfo = { finish_reason: data.finish_reason };
}
generations.push(generation);
}
return { generations };
}
async *_streamResponseChunks(
messages: BaseMessage[],
_options: this["ParsedCallOptions"],
runManager?: CallbackManagerForLLMRun
): AsyncGenerator<ChatGenerationChunk> {
interface ChatFriendliResponse {
choices: {
index: number;
delta: {
role?: FriendliMessageRole;
content?: string;
};
finish_reason: string | null;
}[];
created: number;
}
const response = await this.caller.call(async () =>
fetch(`${this.baseUrl}/v1/chat/completions`, {
method: "POST",
headers: this.constructHeaders(true),
body: this.constructBody(messages, true, _options),
})
);
if (response.status !== 200 ?? !response.body) {
const errorResponse = await response.json();
throw new Error(JSON.stringify(errorResponse));
}
const stream = convertEventStreamToIterableReadableDataStream(
response.body
);
for await (const chunk of stream) {
if (chunk === "[DONE]") break;
const parsedChunk = JSON.parse(chunk) as ChatFriendliResponse;
if (parsedChunk.choices[0].finish_reason === null) {
const generationChunk = new ChatGenerationChunk({
message: _convertDeltaToMessageChunk(parsedChunk.choices[0].delta),
text: parsedChunk.choices[0].delta.content ?? "",
generationInfo: {
finishReason: parsedChunk.choices[0].finish_reason,
},
});
yield generationChunk;
void runManager?.handleLLMNewToken(generationChunk.text ?? "");
}
}
}
}