forked from i-am-alice/3rd-devs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathOpenAIService.ts
111 lines (95 loc) · 3.58 KB
/
OpenAIService.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
import OpenAI from "openai";
import type { ChatCompletionMessageParam } from "openai/resources/chat/completions";
import { createByModelName } from '@microsoft/tiktokenizer';
export class OpenAIService {
private openai: OpenAI;
private tokenizers: Map<string, Awaited<ReturnType<typeof createByModelName>>> = new Map();
private readonly IM_START = "<|im_start|>";
private readonly IM_END = "<|im_end|>";
private readonly IM_SEP = "<|im_sep|>";
constructor() {
this.openai = new OpenAI();
}
private async getTokenizer(modelName: string) {
if (!this.tokenizers.has(modelName)) {
const specialTokens: ReadonlyMap<string, number> = new Map([
[this.IM_START, 100264],
[this.IM_END, 100265],
[this.IM_SEP, 100266],
]);
const tokenizer = await createByModelName(modelName, specialTokens);
this.tokenizers.set(modelName, tokenizer);
}
return this.tokenizers.get(modelName)!;
}
async countTokens(messages: ChatCompletionMessageParam[], model: string = 'gpt-4o'): Promise<number> {
const tokenizer = await this.getTokenizer(model);
let formattedContent = '';
messages.forEach((message) => {
formattedContent += `${this.IM_START}${message.role}${this.IM_SEP}${message.content || ''}${this.IM_END}`;
});
formattedContent += `${this.IM_START}assistant${this.IM_SEP}`;
const tokens = tokenizer.encode(formattedContent, [this.IM_START, this.IM_END, this.IM_SEP]);
return tokens.length;
}
async completion(
messages: ChatCompletionMessageParam[],
model: string = "gpt-4o",
stream: boolean = false,
jsonMode: boolean = false,
maxTokens: number = 4096
): Promise<OpenAI.Chat.Completions.ChatCompletion | AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>> {
try {
const chatCompletion = await this.openai.chat.completions.create({
messages,
model,
...(model !== 'o1-mini' && model !== 'o1-preview' && {
stream,
max_tokens: maxTokens,
response_format: jsonMode ? { type: "json_object" } : { type: "text" }
})
});
if (stream) {
return chatCompletion as AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>;
} else {
return chatCompletion as OpenAI.Chat.Completions.ChatCompletion;
}
} catch (error) {
console.error("Error in OpenAI completion:", error);
throw error;
}
}
async calculateImageTokens(width: number, height: number, detail: 'low' | 'high'): Promise<number> {
let tokenCost = 0;
if (detail === 'low') {
tokenCost += 85;
return tokenCost;
}
const MAX_DIMENSION = 2048;
const SCALE_SIZE = 768;
// Resize to fit within MAX_DIMENSION x MAX_DIMENSION
if (width > MAX_DIMENSION || height > MAX_DIMENSION) {
const aspectRatio = width / height;
if (aspectRatio > 1) {
width = MAX_DIMENSION;
height = Math.round(MAX_DIMENSION / aspectRatio);
} else {
height = MAX_DIMENSION;
width = Math.round(MAX_DIMENSION * aspectRatio);
}
}
// Scale the shortest side to SCALE_SIZE
if (width >= height && height > SCALE_SIZE) {
width = Math.round((SCALE_SIZE / height) * width);
height = SCALE_SIZE;
} else if (height > width && width > SCALE_SIZE) {
height = Math.round((SCALE_SIZE / width) * height);
width = SCALE_SIZE;
}
// Calculate the number of 512px squares
const numSquares = Math.ceil(width / 512) * Math.ceil(height / 512);
// Calculate the token cost
tokenCost += (numSquares * 170) + 85;
return tokenCost;
}
}