-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
minimax.ts
222 lines (192 loc) Β· 7 KB
/
minimax.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
import { getEnvironmentVariable } from "@langchain/core/utils/env";
import { Embeddings, EmbeddingsParams } from "@langchain/core/embeddings";
import { chunkArray } from "@langchain/core/utils/chunk_array";
import { ConfigurationParameters } from "../chat_models/minimax.js";
/**
* Interface for MinimaxEmbeddings parameters. Extends EmbeddingsParams and
* defines additional parameters specific to the MinimaxEmbeddings class.
*/
export interface MinimaxEmbeddingsParams extends EmbeddingsParams {
/** Model name to use */
modelName: string;
/**
* API key to use when making requests. Defaults to the value of
* `MINIMAX_GROUP_ID` environment variable.
*/
minimaxGroupId?: string;
/**
* Secret key to use when making requests. Defaults to the value of
* `MINIMAX_API_KEY` environment variable.
*/
minimaxApiKey?: string;
/**
* The maximum number of documents to embed in a single request. This is
* limited by the Minimax API to a maximum of 4096.
*/
batchSize?: number;
/**
* Whether to strip new lines from the input text. This is recommended by
* Minimax, but may not be suitable for all use cases.
*/
stripNewLines?: boolean;
/**
* The target use-case after generating the vector.
* When using embeddings, the vector of the target content is first generated through the db and stored in the vector database,
* and then the vector of the retrieval text is generated through the query.
* Note: For the parameters of the partial algorithm, we adopted a separate algorithm plan for query and db.
* Therefore, for a paragraph of text, if it is to be used as a retrieval text, it should use the db,
* and if it is used as a retrieval text, it should use the query.
*/
type?: "db" | "query";
}
export interface CreateMinimaxEmbeddingRequest {
/**
* @type {string}
* @memberof CreateMinimaxEmbeddingRequest
*/
model: string;
/**
* Text to generate vector expectation
* @type {CreateEmbeddingRequestInput}
* @memberof CreateMinimaxEmbeddingRequest
*/
texts: string[];
/**
* The target use-case after generating the vector. When using embeddings,
* first generate the vector of the target content through the db and store it in the vector database,
* and then generate the vector of the retrieval text through the query.
* Note: For the parameter of the algorithm, we use the algorithm scheme of query and db separation,
* so a text, if it is to be retrieved as a text, should use the db,
* if it is used as a retrieval text, should use the query.
* @type {string}
* @memberof CreateMinimaxEmbeddingRequest
*/
type: "db" | "query";
}
/**
* Class for generating embeddings using the Minimax API. Extends the
* Embeddings class and implements MinimaxEmbeddingsParams
* @example
* ```typescript
* const embeddings = new MinimaxEmbeddings();
*
* // Embed a single query
* const queryEmbedding = await embeddings.embedQuery("Hello world");
* console.log(queryEmbedding);
*
* // Embed multiple documents
* const documentsEmbedding = await embeddings.embedDocuments([
* "Hello world",
* "Bye bye",
* ]);
* console.log(documentsEmbedding);
* ```
*/
export class MinimaxEmbeddings
extends Embeddings
implements MinimaxEmbeddingsParams
{
modelName = "embo-01";
batchSize = 512;
stripNewLines = true;
minimaxGroupId?: string;
minimaxApiKey?: string;
type: "db" | "query" = "db";
apiUrl: string;
basePath?: string = "https://api.minimax.chat/v1";
headers?: Record<string, string>;
constructor(
fields?: Partial<MinimaxEmbeddingsParams> & {
configuration?: ConfigurationParameters;
}
) {
const fieldsWithDefaults = { maxConcurrency: 2, ...fields };
super(fieldsWithDefaults);
this.minimaxGroupId =
fields?.minimaxGroupId ?? getEnvironmentVariable("MINIMAX_GROUP_ID");
if (!this.minimaxGroupId) {
throw new Error("Minimax GroupID not found");
}
this.minimaxApiKey =
fields?.minimaxApiKey ?? getEnvironmentVariable("MINIMAX_API_KEY");
if (!this.minimaxApiKey) {
throw new Error("Minimax ApiKey not found");
}
this.modelName = fieldsWithDefaults?.modelName ?? this.modelName;
this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize;
this.type = fieldsWithDefaults?.type ?? this.type;
this.stripNewLines =
fieldsWithDefaults?.stripNewLines ?? this.stripNewLines;
this.basePath = fields?.configuration?.basePath ?? this.basePath;
this.apiUrl = `${this.basePath}/embeddings`;
this.headers = fields?.configuration?.headers ?? this.headers;
}
/**
* Method to generate embeddings for an array of documents. Splits the
* documents into batches and makes requests to the Minimax API to generate
* embeddings.
* @param texts Array of documents to generate embeddings for.
* @returns Promise that resolves to a 2D array of embeddings for each document.
*/
async embedDocuments(texts: string[]): Promise<number[][]> {
const batches = chunkArray(
this.stripNewLines ? texts.map((t) => t.replace(/\n/g, " ")) : texts,
this.batchSize
);
const batchRequests = batches.map((batch) =>
this.embeddingWithRetry({
model: this.modelName,
texts: batch,
type: this.type,
})
);
const batchResponses = await Promise.all(batchRequests);
const embeddings: number[][] = [];
for (let i = 0; i < batchResponses.length; i += 1) {
const batch = batches[i];
const { vectors: batchResponse } = batchResponses[i];
for (let j = 0; j < batch.length; j += 1) {
embeddings.push(batchResponse[j]);
}
}
return embeddings;
}
/**
* Method to generate an embedding for a single document. Calls the
* embeddingWithRetry method with the document as the input.
* @param text Document to generate an embedding for.
* @returns Promise that resolves to an embedding for the document.
*/
async embedQuery(text: string): Promise<number[]> {
const { vectors } = await this.embeddingWithRetry({
model: this.modelName,
texts: [this.stripNewLines ? text.replace(/\n/g, " ") : text],
type: this.type,
});
return vectors[0];
}
/**
* Private method to make a request to the Minimax API to generate
* embeddings. Handles the retry logic and returns the response from the
* API.
* @param request Request to send to the Minimax API.
* @returns Promise that resolves to the response from the API.
*/
private async embeddingWithRetry(request: CreateMinimaxEmbeddingRequest) {
const makeCompletionRequest = async () => {
const url = `${this.apiUrl}?GroupId=${this.minimaxGroupId}`;
const response = await fetch(url, {
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${this.minimaxApiKey}`,
...this.headers,
},
body: JSON.stringify(request),
});
const json = await response.json();
return json;
};
return this.caller.call(makeCompletionRequest);
}
}