Skip to content

Commit

Permalink
add function call support
Browse files Browse the repository at this point in the history
  • Loading branch information
YuenSzeHong committed Jan 13, 2025
1 parent de9909b commit c66e172
Showing 1 changed file with 83 additions and 17 deletions.
100 changes: 83 additions & 17 deletions src/api_proxy/worker.mjs
Original file line number Diff line number Diff line change
Expand Up @@ -334,11 +334,44 @@ const transformMessages = async (messages) => {
return { system_instruction, contents };
};

const transformRequest = async (req) => ({
...await transformMessages(req.messages),
safetySettings: safetySettings(req.model),
generationConfig: transformConfig(req),
});
const transformRequest = async (req) => {
const base = {
...await transformMessages(req.messages),
safetySettings: safetySettings(req.model),
generationConfig: transformConfig(req),
};

// Handle functions/tools if present
const tools = req.functions || req.tools;
if (tools) {
// Convert to array if it's a single object
const toolsArray = tools.length ? tools : [tools];

base.tools = [{
functionDeclarations: toolsArray.map(tool => {
// Handle both formats: {type: 'function', function: {name: 'x'}} and {name: 'x'}
const funcDef = tool.type === 'function' ? tool.function : tool;

return {
name: funcDef.name,
description: funcDef.description || "",
parameters: {
type: "object",
properties: {
query: {
type: "string",
description: "The search query to be executed"
}
},
required: ["query"]
}
};
})
}];
}

return base;
};

const generateChatcmplId = () => {
const characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
Expand All @@ -356,15 +389,28 @@ const reasonsMap = { //https://ai.google.dev/api/rest/v1/GenerateContentResponse
// :"function_call",
};
const SEP = "\n\n|>";
const transformCandidates = (key, cand) => ({
index: cand.index || 0, // 0-index is absent in new -002 models response
[key]: {
role: "assistant",
content: cand.content?.parts.map(p => p.text).join(SEP)
},
logprobs: null,
finish_reason: reasonsMap[cand.finishReason] || cand.finishReason,
});
const transformCandidates = (key, cand) => {
const base = {
index: cand.index || 0,
[key]: {
role: "assistant",
content: cand.content?.parts?.map(p => p.text).join(SEP) || null
},
logprobs: null,
finish_reason: reasonsMap[cand.finishReason] || cand.finishReason
};

// Add function_call if present
if (cand.content?.parts?.some(p => p.functionCall)) {
const functionCall = cand.content.parts.find(p => p.functionCall).functionCall;
base[key].function_call = {
name: functionCall.name,
arguments: functionCall.args ? JSON.stringify(functionCall.args) : "{}"
};
}

return base;
};
const transformCandidatesMessage = transformCandidates.bind(null, "message");
const transformCandidatesDelta = transformCandidates.bind(null, "delta");

Expand Down Expand Up @@ -407,21 +453,41 @@ async function parseStreamFlush(controller) {

function transformResponseStream(data, stop, first) {
const item = transformCandidatesDelta(data.candidates[0]);
if (stop) { item.delta = {}; } else { item.finish_reason = null; }
if (first) { item.delta.content = ""; } else { delete item.delta.role; }
if (stop) {
item.delta = {};
} else {
item.finish_reason = null;
}

if (first) {
item.delta.content = "";
// Add function_call structure if present
if (data.candidates[0].content?.parts?.some(p => p.functionCall)) {
const functionCall = data.candidates[0].content.parts.find(p => p.functionCall).functionCall;
item.delta.function_call = {
name: functionCall.name,
arguments: functionCall.args ? JSON.stringify(functionCall.args) : "{}"
};
}
} else {
delete item.delta.role;
}

const output = {
id: this.id,
choices: [item],
created: Math.floor(Date.now() / 1000),
model: this.model,
//system_fingerprint: "fp_69829325d0",
object: "chat.completion.chunk",
};

if (data.usageMetadata && this.streamIncludeUsage) {
output.usage = stop ? transformUsage(data.usageMetadata) : null;
}

return "data: " + JSON.stringify(output) + delimiter;
}

const delimiter = "\n\n";
async function toOpenAiStream(chunk, controller) {
const transform = transformResponseStream.bind(this);
Expand Down

0 comments on commit c66e172

Please sign in to comment.