-
Notifications
You must be signed in to change notification settings - Fork 314
/
open_ai.ts
60 lines (51 loc) · 1.63 KB
/
open_ai.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
import _ from "lodash";
import OpenAI, { ClientOptions } from "openai";
export class AzureOpenAI extends OpenAI {
protected override authHeaders() {
return { "api-key": this.apiKey };
}
}
export class OpenAISession {
openai: OpenAI;
constructor(options: ClientOptions & { azure?: boolean } = {}) {
if (!options.apiKey) {
if (typeof process !== undefined) {
options.apiKey = process.env.OPENAI_API_KEY;
}
}
if (!options.apiKey) {
throw new Error("Set OpenAI Key in OPENAI_API_KEY env variable"); // Overriding OpenAI package's error message
}
if (options.azure) {
this.openai = new AzureOpenAI(options);
} else {
this.openai = new OpenAI({
...options,
// defaultHeaders: { "OpenAI-Beta": "assistants=v1" },
});
}
}
}
// I'm not 100% sure this is necessary vs. just starting a new session
// every time we make a call. They say they try to reuse connections
// so in theory this is more efficient, but we should test it in the future.
let defaultOpenAISession: { session: OpenAISession; options: ClientOptions }[] =
[];
/**
* Get a session for the OpenAI API. If one already exists with the same options,
* it will be returned. Otherwise, a new session will be created.
* @param options
* @returns
*/
export function getOpenAISession(
options: ClientOptions & { azure?: boolean } = {},
) {
let session = defaultOpenAISession.find((session) => {
return _.isEqual(session.options, options);
})?.session;
if (!session) {
session = new OpenAISession(options);
defaultOpenAISession.push({ session, options });
}
return session;
}