Skip to content

Commit f0d35e2

Browse files
authored
feat: vertex ai openai adapter (#6736)
* feat: vertex ai openai adapter * fix: vertex ai - missing mistral key * feat: support tools for vertexAI gemini and anthropic * feat: vertex openai adapter working for gemini * fix: show errors in vertex ai * fix: filter out model from vertexai anthropic body * feat: vertexai toolSupport update * chore: remove force openai adapter usage for vertex * fix: prettier
1 parent 70302e2 commit f0d35e2

File tree

12 files changed

+980
-238
lines changed

12 files changed

+980
-238
lines changed

core/llm/llms/Anthropic.ts

Lines changed: 53 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -172,52 +172,10 @@ class Anthropic extends BaseLLM {
172172
}
173173
}
174174

175-
protected async *_streamChat(
176-
messages: ChatMessage[],
177-
signal: AbortSignal,
178-
options: CompletionOptions,
175+
async *handleResponse(
176+
response: any,
177+
stream: boolean | undefined,
179178
): AsyncGenerator<ChatMessage> {
180-
if (!this.apiKey || this.apiKey === "") {
181-
throw new Error(
182-
"Request not sent. You have an Anthropic model configured in your config.json, but the API key is not set.",
183-
);
184-
}
185-
186-
const systemMessage = stripImages(
187-
messages.filter((m) => m.role === "system")[0]?.content ?? "",
188-
);
189-
const shouldCacheSystemMessage = !!(
190-
this.cacheBehavior?.cacheSystemMessage && systemMessage
191-
);
192-
193-
const msgs = this.convertMessages(messages);
194-
const response = await this.fetch(new URL("messages", this.apiBase), {
195-
method: "POST",
196-
headers: {
197-
"Content-Type": "application/json",
198-
Accept: "application/json",
199-
"anthropic-version": "2023-06-01",
200-
"x-api-key": this.apiKey as string,
201-
...(shouldCacheSystemMessage || this.cacheBehavior?.cacheConversation
202-
? { "anthropic-beta": "prompt-caching-2024-07-31" }
203-
: {}),
204-
},
205-
body: JSON.stringify({
206-
...this.convertArgs(options),
207-
messages: msgs,
208-
system: shouldCacheSystemMessage
209-
? [
210-
{
211-
type: "text",
212-
text: systemMessage,
213-
cache_control: { type: "ephemeral" },
214-
},
215-
]
216-
: systemMessage,
217-
}),
218-
signal,
219-
});
220-
221179
if (response.status === 499) {
222180
return; // Aborted by user
223181
}
@@ -237,7 +195,7 @@ class Anthropic extends BaseLLM {
237195
);
238196
}
239197

240-
if (options.stream === false) {
198+
if (stream === false) {
241199
const data = await response.json();
242200
const cost = data.usage
243201
? {
@@ -348,6 +306,55 @@ class Anthropic extends BaseLLM {
348306
usage,
349307
};
350308
}
309+
310+
protected async *_streamChat(
311+
messages: ChatMessage[],
312+
signal: AbortSignal,
313+
options: CompletionOptions,
314+
): AsyncGenerator<ChatMessage> {
315+
if (!this.apiKey || this.apiKey === "") {
316+
throw new Error(
317+
"Request not sent. You have an Anthropic model configured in your config.json, but the API key is not set.",
318+
);
319+
}
320+
321+
const systemMessage = stripImages(
322+
messages.filter((m) => m.role === "system")[0]?.content ?? "",
323+
);
324+
const shouldCacheSystemMessage = !!(
325+
this.cacheBehavior?.cacheSystemMessage && systemMessage
326+
);
327+
328+
const msgs = this.convertMessages(messages);
329+
const response = await this.fetch(new URL("messages", this.apiBase), {
330+
method: "POST",
331+
headers: {
332+
"Content-Type": "application/json",
333+
Accept: "application/json",
334+
"anthropic-version": "2023-06-01",
335+
"x-api-key": this.apiKey as string,
336+
...(shouldCacheSystemMessage || this.cacheBehavior?.cacheConversation
337+
? { "anthropic-beta": "prompt-caching-2024-07-31" }
338+
: {}),
339+
},
340+
body: JSON.stringify({
341+
...this.convertArgs(options),
342+
messages: msgs,
343+
system: shouldCacheSystemMessage
344+
? [
345+
{
346+
type: "text",
347+
text: systemMessage,
348+
cache_control: { type: "ephemeral" },
349+
},
350+
]
351+
: systemMessage,
352+
}),
353+
signal,
354+
});
355+
356+
yield* this.handleResponse(response, options.stream);
357+
}
351358
}
352359

353360
export default Anthropic;

core/llm/llms/Gemini.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ class Gemini extends BaseLLM {
200200
messages: ChatMessage[],
201201
options: CompletionOptions,
202202
isV1API: boolean,
203+
includeToolIds: boolean,
203204
): GeminiChatRequestBody {
204205
const toolCallIdToNameMap = new Map<string, string>();
205206
messages.forEach((msg) => {
@@ -231,7 +232,7 @@ class Gemini extends BaseLLM {
231232
parts: [
232233
{
233234
functionResponse: {
234-
id: msg.toolCallId,
235+
id: includeToolIds ? msg.toolCallId : undefined,
235236
name: functionName || "unknown",
236237
response: {
237238
output: msg.content, // "output" key is opinionated - not all functions will output objects
@@ -406,7 +407,7 @@ class Gemini extends BaseLLM {
406407
const isV1API = !!this.apiBase?.includes("/v1/");
407408

408409
// Convert chat messages to contents
409-
const body = this.prepareBody(messages, options, isV1API);
410+
const body = this.prepareBody(messages, options, isV1API, true);
410411

411412
const response = await this.fetch(apiURL, {
412413
method: "POST",

core/llm/llms/Mistral.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ class Mistral extends OpenAI {
6767
return Mistral.modelConversion[model] ?? model;
6868
}
6969

70-
protected _convertArgs(options: any, messages: ChatMessage[]) {
70+
_convertArgs(options: any, messages: ChatMessage[]) {
7171
const finalOptions = super._convertArgs(options, messages);
7272

7373
const lastMessage = finalOptions.messages[finalOptions.messages.length - 1];

core/llm/llms/VertexAI.ts

Lines changed: 9 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -254,24 +254,7 @@ class VertexAI extends BaseLLM {
254254
signal,
255255
});
256256

257-
if (response.status === 499) {
258-
return; // Aborted by user
259-
}
260-
261-
if (options.stream === false) {
262-
const data = await response.json();
263-
yield { role: "assistant", content: data.content[0].text };
264-
return;
265-
}
266-
267-
for await (const value of streamSse(response)) {
268-
if (value.type === "message_start") {
269-
console.log(value);
270-
}
271-
if (value.delta?.text) {
272-
yield { role: "assistant", content: value.delta.text };
273-
}
274-
}
257+
yield* this.anthropicInstance.handleResponse(response, options.stream);
275258
}
276259

277260
// Gemini
@@ -285,17 +268,19 @@ class VertexAI extends BaseLLM {
285268
this.apiBase,
286269
);
287270

288-
const body = this.geminiInstance.prepareBody(messages, options, false);
271+
// For some reason gemini through vertex does not support ids in functionResponses yet
272+
const body = this.geminiInstance.prepareBody(
273+
messages,
274+
options,
275+
false,
276+
false,
277+
);
289278
const response = await this.fetch(apiURL, {
290279
method: "POST",
291280
body: JSON.stringify(body),
292281
signal,
293282
});
294-
for await (const message of this.geminiInstance.processGeminiResponse(
295-
streamResponse(response),
296-
)) {
297-
yield message;
298-
}
283+
yield* this.geminiInstance.processGeminiResponse(streamResponse(response));
299284
}
300285

301286
private async *streamChatBison(

core/llm/toolSupport.ts

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,11 +83,12 @@ export const PROVIDER_TOOL_SUPPORT: Record<string, (model: string) => boolean> =
8383
return model.toLowerCase().includes("gemini");
8484
},
8585
vertexai: (model) => {
86+
const lowerCaseModel = model.toLowerCase();
8687
// All gemini models except flash 2.0 lite support function calling
87-
return (
88-
model.toLowerCase().includes("gemini") &&
89-
!model.toLowerCase().includes("lite")
90-
);
88+
if (lowerCaseModel.includes("lite")) {
89+
return false;
90+
}
91+
return ["claude", "gemini"].some((val) => lowerCaseModel.includes(val));
9192
},
9293
bedrock: (model) => {
9394
if (

0 commit comments

Comments
 (0)