dev: gemini function calling

This commit is contained in:
KernelDeimos
2025-02-19 13:16:27 -05:00
parent d8c49d8b2d
commit 9d181741de
3 changed files with 82 additions and 2 deletions
@@ -3,6 +3,7 @@ const { GoogleGenerativeAI } = require('@google/generative-ai');
const GeminiSquareHole = require("./lib/GeminiSquareHole");
const { TypedValue } = require("../../services/drivers/meta/Runtime");
const putility = require("@heyputer/putility");
const FunctionCalling = require("./lib/FunctionCalling");
class GeminiService extends BaseService {
async _init () {
@@ -31,9 +32,12 @@ class GeminiService extends BaseService {
},
async complete ({ messages, stream, model, tools }) {
tools = FunctionCalling.make_gemini_tools(tools);
const genAI = new GoogleGenerativeAI(this.config.apiKey);
const genModel = genAI.getGenerativeModel({
model: model ?? 'gemini-2.0-flash',
tools,
});
messages = await GeminiSquareHole.process_input_messages(messages);
@@ -41,7 +45,9 @@ class GeminiService extends BaseService {
// History is separate, so the last message gets special treatment.
const last_message = messages.pop();
const last_message_parts = last_message.parts.map(
part => typeof part === 'string' ? part : part.text
part => typeof part === 'string' ? part :
typeof part.text === 'string' ? part.text :
part
);
const chat = genModel.startChat({
@@ -90,4 +90,16 @@ module.exports = class FunctionCalling {
};
});
}
static make_gemini_tools (tools) {
return [
{
function_declarations: tools.map(t => {
const tool = t.function;
delete tool.parameters.additionalProperties;
return tool;
})
}
];
}
}
@@ -3,6 +3,8 @@
* but Google's AI API defies all the established conventions
* so it made sense to defy them here as well.
*/
const crypto = require('crypto');
module.exports = class GeminiSquareHole {
static process_input_messages = async (messages) => {
messages = messages.slice();
@@ -14,6 +16,35 @@ module.exports = class GeminiSquareHole {
if ( msg.role === 'assistant' ) {
msg.role = 'model';
}
for ( let i=0 ; i < msg.parts.length ; i++ ) {
const part = msg.parts[i];
console.log('what the part is', part);
if ( part.type === 'tool_use' ) {
msg.parts[i] = {
functionCall: {
name: part.id,
args: part.input,
},
};
}
if ( part.type === 'tool_result' ) {
msg.parts[i] = {
functionResponse: {
name: part.tool_use_id,
response: {
name: part.tool_use_id,
content: part.content,
},
},
};
}
if ( part.type === 'text' ) {
msg.parts[i] = {
text: part.text,
};
}
}
}
return messages;
@@ -46,7 +77,12 @@ module.exports = class GeminiSquareHole {
usage_promise,
}) => async ({ chatStream }) => {
const message = chatStream.message();
let textblock = message.contentBlock({ type: 'text' });
let toolblock = null;
let mode = 'text';
let last_usage = null;
for await ( const chunk of stream ) {
// This is spread across several lines so that the stack trace
@@ -56,6 +92,31 @@ module.exports = class GeminiSquareHole {
const content = candidate.content;
const parts = content.parts;
for ( const part of parts ) {
if ( part.functionCall ) {
if ( mode === 'text' ) {
mode = 'tool';
textblock.end();
}
toolblock = message.contentBlock({
type: 'tool_use',
id: part.functionCall.name,
name: part.functionCall.name,
});
toolblock.addPartialJSON(JSON.stringify(
part.functionCall.args,
));
continue;
}
if ( mode === 'tool' ) {
mode = 'text';
toolblock.end();
textblock = message.contentBlock({ type: 'text' });
}
// assume text as default
const text = part.text;
textblock.addText(text);
}
@@ -65,7 +126,8 @@ module.exports = class GeminiSquareHole {
usage_promise.resolve(last_usage);
textblock.end();
if ( mode === 'text' ) textblock.end();
if ( mode === 'tool' ) toolblock.end();
message.end();
chatStream.end();
}