Skip to content

Commit

Permalink
Merge pull request #5581 from ConnectAI-E/feature/gemini-functioncall
Browse files Browse the repository at this point in the history
google gemini support function call
  • Loading branch information
lloydzhou authored Oct 10, 2024
2 parents ba58018 + 4ae34ea commit c5074f0
Show file tree
Hide file tree
Showing 3 changed files with 86 additions and 112 deletions.
194 changes: 82 additions & 112 deletions app/client/platforms/google.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,25 @@ import {
LLMUsage,
SpeechOptions,
} from "../api";
import { useAccessStore, useAppConfig, useChatStore } from "@/app/store";
import {
useAccessStore,
useAppConfig,
useChatStore,
usePluginStore,
ChatMessageTool,
} from "@/app/store";
import { stream } from "@/app/utils/chat";
import { getClientConfig } from "@/app/config/client";
import { GEMINI_BASE_URL } from "@/app/constant";
import Locale from "../../locales";
import {
EventStreamContentType,
fetchEventSource,
} from "@fortaine/fetch-event-source";
import { prettyObject } from "@/app/utils/format";

import {
getMessageTextContent,
getMessageImages,
isVisionModel,
} from "@/app/utils";
import { preProcessImageContent } from "@/app/utils/chat";
import { nanoid } from "nanoid";
import { RequestPayload } from "./openai";
import { fetch } from "@/app/utils/stream";

export class GeminiProApi implements LLMApi {
Expand Down Expand Up @@ -178,115 +182,81 @@ export class GeminiProApi implements LLMApi {
);

if (shouldStream) {
let responseText = "";
let remainText = "";
let finished = false;

const finish = () => {
if (!finished) {
finished = true;
options.onFinish(responseText + remainText);
}
};

// animate response to make it looks smooth
function animateResponseText() {
if (finished || controller.signal.aborted) {
responseText += remainText;
finish();
return;
}

if (remainText.length > 0) {
const fetchCount = Math.max(1, Math.round(remainText.length / 60));
const fetchText = remainText.slice(0, fetchCount);
responseText += fetchText;
remainText = remainText.slice(fetchCount);
options.onUpdate?.(responseText, fetchText);
}

requestAnimationFrame(animateResponseText);
}

// start animaion
animateResponseText();

controller.signal.onabort = finish;

fetchEventSource(chatPath, {
fetch: fetch as any,
...chatPayload,
async onopen(res) {
clearTimeout(requestTimeoutId);
const contentType = res.headers.get("content-type");
console.log(
"[Gemini] request response content type: ",
contentType,
);

if (contentType?.startsWith("text/plain")) {
responseText = await res.clone().text();
return finish();
}

if (
!res.ok ||
!res.headers
.get("content-type")
?.startsWith(EventStreamContentType) ||
res.status !== 200
) {
const responseTexts = [responseText];
let extraInfo = await res.clone().text();
try {
const resJson = await res.clone().json();
extraInfo = prettyObject(resJson);
} catch {}

if (res.status === 401) {
responseTexts.push(Locale.Error.Unauthorized);
}

if (extraInfo) {
responseTexts.push(extraInfo);
}

responseText = responseTexts.join("\n\n");

return finish();
}
},
onmessage(msg) {
if (msg.data === "[DONE]" || finished) {
return finish();
}
const text = msg.data;
try {
const json = JSON.parse(text);
const delta = apiClient.extractMessage(json);

if (delta) {
remainText += delta;
}
const [tools, funcs] = usePluginStore
.getState()
.getAsTools(
useChatStore.getState().currentSession().mask?.plugin || [],
);
return stream(
chatPath,
requestPayload,
getHeaders(),
// @ts-ignore
[{ functionDeclarations: tools.map((tool) => tool.function) }],
funcs,
controller,
// parseSSE
(text: string, runTools: ChatMessageTool[]) => {
// console.log("parseSSE", text, runTools);
const chunkJson = JSON.parse(text);

const blockReason = json?.promptFeedback?.blockReason;
if (blockReason) {
// being blocked
console.log(`[Google] [Safety Ratings] result:`, blockReason);
}
} catch (e) {
console.error("[Request] parse error", text, msg);
const functionCall = chunkJson?.candidates
?.at(0)
?.content.parts.at(0)?.functionCall;
if (functionCall) {
const { name, args } = functionCall;
runTools.push({
id: nanoid(),
type: "function",
function: {
name,
arguments: JSON.stringify(args), // utils.chat call function, using JSON.parse
},
});
}
return chunkJson?.candidates?.at(0)?.content.parts.at(0)?.text;
},
onclose() {
finish();
},
onerror(e) {
options.onError?.(e);
throw e;
// processToolMessage, include tool_calls message and tool call results
(
requestPayload: RequestPayload,
toolCallMessage: any,
toolCallResult: any[],
) => {
// @ts-ignore
requestPayload?.contents?.splice(
// @ts-ignore
requestPayload?.contents?.length,
0,
{
role: "model",
parts: toolCallMessage.tool_calls.map(
(tool: ChatMessageTool) => ({
functionCall: {
name: tool?.function?.name,
args: JSON.parse(tool?.function?.arguments as string),
},
}),
),
},
// @ts-ignore
...toolCallResult.map((result) => ({
role: "function",
parts: [
{
functionResponse: {
name: result.name,
response: {
name: result.name,
content: result.content, // TODO just text content...
},
},
},
],
})),
);
},
openWhenHidden: true,
});
options,
);
} else {
const res = await fetch(chatPath, chatPayload);
clearTimeout(requestTimeoutId);
Expand Down
3 changes: 3 additions & 0 deletions app/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,9 @@ export function showPlugins(provider: ServiceProvider, model: string) {
if (provider == ServiceProvider.Anthropic && !model.includes("claude-2")) {
return true;
}
if (provider == ServiceProvider.Google && !model.includes("vision")) {
return true;
}
return false;
}

Expand Down
1 change: 1 addition & 0 deletions app/utils/chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,7 @@ export function stream(
return e.toString();
})
.then((content) => ({
name: tool.function.name,
role: "tool",
content,
tool_call_id: tool.id,
Expand Down

0 comments on commit c5074f0

Please sign in to comment.