[VertexAI] Add support for tools parameter (#1065)

* [VertexAI] Add support for tools parameter * Simplify tools parameter parsing and add support for passing parameters in model --------- Co-authored-by: Nathan Sarrazin <[email protected]>
huggingface · May 3, 2024 · 50febad · 50febad
1 parent 5703512
commit 50febad
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -619,7 +619,12 @@ MODELS=`[
 
           // Optional
           "safetyThreshold": "BLOCK_MEDIUM_AND_ABOVE",
-          "apiEndpoint": "", // alternative api endpoint url
+          "apiEndpoint": "", // alternative api endpoint url,
+          "tools": [{
+            "googleSearchRetrieval": {
+              "disableAttribution": true
+            }
+          }]
        }]
      },
 ]`

diff --git a/src/lib/server/endpoints/google/endpointVertex.ts b/src/lib/server/endpoints/google/endpointVertex.ts
@@ -26,10 +26,11 @@ export const endpointVertexParametersSchema = z.object({
 			HarmBlockThreshold.BLOCK_ONLY_HIGH,
 		])
 		.optional(),
+	tools: z.array(z.any()),
 });
 
 export function endpointVertex(input: z.input<typeof endpointVertexParametersSchema>): Endpoint {
-	const { project, location, model, apiEndpoint, safetyThreshold } =
+	const { project, location, model, apiEndpoint, safetyThreshold, tools } =
 		endpointVertexParametersSchema.parse(input);
 
 	const vertex_ai = new VertexAI({
@@ -39,6 +40,8 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
 	});
 
 	return async ({ messages, preprompt, generateSettings }) => {
+		const parameters = { ...model.parameters, ...generateSettings };
+
 		const generativeModel = vertex_ai.getGenerativeModel({
 			model: model.id ?? model.name,
 			safetySettings: safetyThreshold
@@ -66,10 +69,11 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
 				  ]
 				: undefined,
 			generationConfig: {
-				maxOutputTokens: generateSettings?.max_new_tokens ?? 4096,
-				stopSequences: generateSettings?.stop,
-				temperature: generateSettings?.temperature ?? 1,
+				maxOutputTokens: parameters?.max_new_tokens ?? 4096,
+				stopSequences: parameters?.stop,
+				temperature: parameters?.temperature ?? 1,
 			},
+			tools,
 		});
 
 		// Preprompt is the same as the first system message.