📦 NEW: Google tool call support (#53)

Author: msaaddev

packages/baseai/src/dev/llms/call-google.ts

@@ -5,6 +5,7 @@ import { GOOGLE } from '../data/models';
 import { applyJsonModeIfEnabledForGoogle, handleLlmError } from './utils';
 import type { ModelParams } from 'types/providers';
 import type { Message } from 'types/pipe';
+import { addToolsToParams } from '../utils/add-tools-to-params';
 
 export async function callGoogle({
 	pipe,
@@ -19,6 +20,7 @@ export async function callGoogle({
 }) {
 	try {
 		const modelParams = buildModelParams(pipe, stream, messages);
+		addToolsToParams(modelParams, pipe);
 
 		// Transform params according to provider's format
 		const transformedRequestParams = transformToProviderRequest({

packages/baseai/src/dev/providers/google/chatComplete.ts

@@ -1,12 +1,15 @@
 import {
 	generateErrorResponse,
-	generateInvalidProviderResponseError
+	generateInvalidProviderResponseError,
+	getMimeType
 } from '../utils';
 import { GOOGLE } from '@/dev/data/models';
+import type { ToolCall, ToolChoice } from 'types/pipe';
 import type {
 	ChatCompletionResponse,
 	ContentType,
 	ErrorResponse,
+	MessageRole,
 	ModelParams,
 	ProviderConfig,
 	ProviderMessage
@@ -32,6 +35,76 @@ const transformGenerationConfig = (params: ModelParams) => {
 	return generationConfig;
 };
 
+export type GoogleMessageRole = 'user' | 'model' | 'function';
+
+interface GoogleFunctionCallMessagePart {
+	functionCall: GoogleGenerateFunctionCall;
+}
+
+interface GoogleFunctionResponseMessagePart {
+	functionResponse: {
+		name: string;
+		response: {
+			name?: string;
+			content: string;
+		};
+	};
+}
+
+type GoogleMessagePart =
+	| GoogleFunctionCallMessagePart
+	| GoogleFunctionResponseMessagePart
+	| { text: string };
+
+export interface GoogleMessage {
+	role: GoogleMessageRole;
+	parts: GoogleMessagePart[];
+}
+
+export interface GoogleToolConfig {
+	function_calling_config: {
+		mode: GoogleToolChoiceType | undefined;
+		allowed_function_names?: string[];
+	};
+}
+
+export const transformOpenAIRoleToGoogleRole = (
+	role: MessageRole
+): GoogleMessageRole => {
+	switch (role) {
+		case 'assistant':
+			return 'model';
+		case 'tool':
+			return 'function';
+		// Not all gemini models support system role
+		case 'system':
+			return 'user';
+		// user is the default role
+		default:
+			return role;
+	}
+};
+
+type GoogleToolChoiceType = 'AUTO' | 'ANY' | 'NONE';
+
+export const transformToolChoiceForGemini = (
+	tool_choice: ToolChoice
+): GoogleToolChoiceType | undefined => {
+	if (typeof tool_choice === 'object' && tool_choice.type === 'function')
+		return 'ANY';
+	if (typeof tool_choice === 'string') {
+		switch (tool_choice) {
+			case 'auto':
+				return 'AUTO';
+			case 'none':
+				return 'NONE';
+			case 'required':
+				return 'ANY';
+		}
+	}
+	return undefined;
+};
+
 export const GoogleChatCompleteConfig: ProviderConfig = {
 	model: {
 		param: 'model',
@@ -42,36 +115,100 @@ export const GoogleChatCompleteConfig: ProviderConfig = {
 		param: 'contents',
 		default: '',
 		transform: (params: ModelParams) => {
-			const messages: { role: string; parts: { text: string }[] }[] = [];
+			const messages: GoogleMessage[] = [];
+			let lastRole: GoogleMessageRole | undefined;
 
 			params.messages?.forEach((message: ProviderMessage) => {
-				const role = message.role === 'assistant' ? 'model' : 'user';
+				const role = transformOpenAIRoleToGoogleRole(message.role);
 				let parts = [];
-				if (typeof message.content === 'string') {
+
+				if (message.role === 'assistant' && message.tool_calls) {
+					message.tool_calls.forEach((tool_call: ToolCall) => {
+						parts.push({
+							functionCall: {
+								name: tool_call.function.name,
+								args: JSON.parse(tool_call.function.arguments)
+							}
+						});
+					});
+				} else if (
+					message.role === 'tool' &&
+					typeof message.content === 'string'
+				) {
 					parts.push({
-						text: message.content
+						functionResponse: {
+							name: message.name ?? 'lb-random-tool-name',
+							response: {
+								content: message.content
+							}
+						}
 					});
-				}
-
-				if (message.content && typeof message.content === 'object') {
+				} else if (
+					message.content &&
+					typeof message.content === 'object'
+				) {
 					message.content.forEach((c: ContentType) => {
 						if (c.type === 'text') {
 							parts.push({
 								text: c.text
 							});
 						}
 						if (c.type === 'image_url') {
-							parts.push({
-								inlineData: {
-									mimeType: 'image/jpeg',
-									data: c.image_url?.url
-								}
-							});
+							const { url } = c.image_url || {};
+							if (!url) return;
+
+							// Handle different types of image URLs
+							if (url.startsWith('data:')) {
+								const [mimeTypeWithPrefix, base64Image] =
+									url.split(';base64,');
+								const mimeType =
+									mimeTypeWithPrefix.split(':')[1];
+
+								parts.push({
+									inlineData: {
+										mimeType: mimeType,
+										data: base64Image
+									}
+								});
+							} else if (
+								url.startsWith('gs://') ||
+								url.startsWith('https://') ||
+								url.startsWith('http://')
+							) {
+								parts.push({
+									fileData: {
+										mimeType: getMimeType(url),
+										fileUri: url
+									}
+								});
+							} else {
+								parts.push({
+									inlineData: {
+										mimeType: 'image/jpeg',
+										data: c.image_url?.url
+									}
+								});
+							}
 						}
 					});
+				} else if (typeof message.content === 'string') {
+					parts.push({
+						text: message.content
+					});
 				}
 
-				messages.push({ role, parts });
+				// Combine consecutive messages if they are from the same role
+				// This takes care of the "Please ensure that multiturn requests alternate between user and model.
+				// Also possible fix for "Please ensure that function call turn comes immediately after a user turn or after a function response turn." in parallel tool calls
+				const shouldCombineMessages =
+					lastRole === role && !params.model?.includes('vision');
+
+				if (shouldCombineMessages) {
+					messages[messages.length - 1].parts.push(...parts);
+				} else {
+					messages.push({ role, parts });
+				}
+				lastRole = role;
 			});
 			return messages;
 		}
@@ -108,6 +245,36 @@ export const GoogleChatCompleteConfig: ProviderConfig = {
 			});
 			return [{ functionDeclarations }];
 		}
+	},
+	tool_choice: {
+		param: 'tool_config',
+		default: '',
+		transform: (params: ModelParams) => {
+			if (params.tool_choice) {
+				const allowedFunctionNames: string[] = [];
+				// If tool_choice is an object and type is function, add the function name to allowedFunctionNames
+				if (
+					typeof params.tool_choice === 'object' &&
+					params.tool_choice.type === 'function'
+				) {
+					allowedFunctionNames.push(params.tool_choice.function.name);
+				}
+				const toolConfig: GoogleToolConfig = {
+					function_calling_config: {
+						mode: transformToolChoiceForGemini(params.tool_choice)
+					}
+				};
+				// TODO: @msaaddev I think we can't have more than one function in tool_choice
+				// but this will also handle the case if we have more than one function in tool_choice
+
+				// If tool_choice has functions, add the function names to allowedFunctionNames
+				if (allowedFunctionNames.length > 0) {
+					toolConfig.function_calling_config.allowed_function_names =
+						allowedFunctionNames;
+				}
+				return toolConfig;
+			}
+		}
 	}
 };
 
@@ -146,6 +313,11 @@ interface GoogleGenerateContentResponse {
 			probability: string;
 		}[];
 	};
+	usageMetadata: {
+		promptTokenCount: number;
+		candidatesTokenCount: number;
+		totalTokenCount: number;
+	};
 }
 
 export const GoogleChatCompleteResponseTransform: (
@@ -170,7 +342,6 @@ export const GoogleChatCompleteResponseTransform: (
 			GOOGLE
 		);
 	}
-
 	if ('candidates' in response) {
 		return {
 			id: crypto.randomUUID(),
@@ -179,7 +350,7 @@ export const GoogleChatCompleteResponseTransform: (
 			model: 'Unknown',
 			provider: GOOGLE,
 			choices:
-				response.candidates?.map((generation, index) => {
+				response.candidates?.map(generation => {
 					// In blocking mode: Google AI does not return content if response > max output tokens param
 					// Test it by asking a big response while keeping maxtokens low ~ 50
 					if (
@@ -203,28 +374,34 @@ export const GoogleChatCompleteResponseTransform: (
 					} else if (generation.content?.parts[0]?.functionCall) {
 						message = {
 							role: 'assistant',
-							tool_calls: [
-								{
-									id: crypto.randomUUID(),
-									type: 'function',
-									function: {
-										name: generation.content.parts[0]
-											?.functionCall.name,
-										arguments: JSON.stringify(
-											generation.content.parts[0]
-												?.functionCall.args
-										)
-									}
+							content: null,
+							tool_calls: generation.content.parts.map(part => {
+								if (part.functionCall) {
+									return {
+										id: crypto.randomUUID(),
+										type: 'function',
+										function: {
+											name: part.functionCall.name,
+											arguments: JSON.stringify(
+												part.functionCall.args
+											)
+										}
+									};
 								}
-							]
+							})
 						};
 					}
 					return {
 						message: message,
 						index: generation.index,
 						finish_reason: generation.finishReason
 					};
-				}) ?? []
+				}) ?? [],
+			usage: {
+				prompt_tokens: response.usageMetadata.promptTokenCount,
+				completion_tokens: response.usageMetadata.candidatesTokenCount,
+				total_tokens: response.usageMetadata.totalTokenCount
+			}
 		};
 	}
 
@@ -262,7 +439,7 @@ export const GoogleChatCompleteStreamChunkTransform: (
 			model: '',
 			provider: 'google',
 			choices:
-				parsedChunk.candidates?.map((generation, index) => {
+				parsedChunk.candidates?.map(generation => {
 					let message: ProviderMessage = {
 						role: 'assistant',
 						content: ''
@@ -275,21 +452,23 @@ export const GoogleChatCompleteStreamChunkTransform: (
 					} else if (generation.content.parts[0]?.functionCall) {
 						message = {
 							role: 'assistant',
-							tool_calls: [
-								{
-									id: crypto.randomUUID(),
-									type: 'function',
-									index: 0,
-									function: {
-										name: generation.content.parts[0]
-											?.functionCall.name,
-										arguments: JSON.stringify(
-											generation.content.parts[0]
-												?.functionCall.args
-										)
+							tool_calls: generation.content.parts.map(
+								(part, idx) => {
+									if (part.functionCall) {
+										return {
+											index: idx,
+											id: crypto.randomUUID(),
+											type: 'function',
+											function: {
+												name: part.functionCall.name,
+												arguments: JSON.stringify(
+													part.functionCall.args
+												)
+											}
+										};
 									}
 								}
-							]
+							)
 						};
 					}
 					return {