Spaces:
Paused
Paused
[VertexAI] Add support for tools parameter (#1065)
Browse files* [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]>
README.md
CHANGED
|
@@ -619,7 +619,12 @@ MODELS=`[
|
|
| 619 |
|
| 620 |
// Optional
|
| 621 |
"safetyThreshold": "BLOCK_MEDIUM_AND_ABOVE",
|
| 622 |
-
"apiEndpoint": "", // alternative api endpoint url
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 623 |
}]
|
| 624 |
},
|
| 625 |
]`
|
|
|
|
| 619 |
|
| 620 |
// Optional
|
| 621 |
"safetyThreshold": "BLOCK_MEDIUM_AND_ABOVE",
|
| 622 |
+
"apiEndpoint": "", // alternative api endpoint url,
|
| 623 |
+
"tools": [{
|
| 624 |
+
"googleSearchRetrieval": {
|
| 625 |
+
"disableAttribution": true
|
| 626 |
+
}
|
| 627 |
+
}]
|
| 628 |
}]
|
| 629 |
},
|
| 630 |
]`
|
src/lib/server/endpoints/google/endpointVertex.ts
CHANGED
|
@@ -26,10 +26,11 @@ export const endpointVertexParametersSchema = z.object({
|
|
| 26 |
HarmBlockThreshold.BLOCK_ONLY_HIGH,
|
| 27 |
])
|
| 28 |
.optional(),
|
|
|
|
| 29 |
});
|
| 30 |
|
| 31 |
export function endpointVertex(input: z.input<typeof endpointVertexParametersSchema>): Endpoint {
|
| 32 |
-
const { project, location, model, apiEndpoint, safetyThreshold } =
|
| 33 |
endpointVertexParametersSchema.parse(input);
|
| 34 |
|
| 35 |
const vertex_ai = new VertexAI({
|
|
@@ -39,6 +40,8 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
|
|
| 39 |
});
|
| 40 |
|
| 41 |
return async ({ messages, preprompt, generateSettings }) => {
|
|
|
|
|
|
|
| 42 |
const generativeModel = vertex_ai.getGenerativeModel({
|
| 43 |
model: model.id ?? model.name,
|
| 44 |
safetySettings: safetyThreshold
|
|
@@ -66,10 +69,11 @@ export function endpointVertex(input: z.input<typeof endpointVertexParametersSch
|
|
| 66 |
]
|
| 67 |
: undefined,
|
| 68 |
generationConfig: {
|
| 69 |
-
maxOutputTokens:
|
| 70 |
-
stopSequences:
|
| 71 |
-
temperature:
|
| 72 |
},
|
|
|
|
| 73 |
});
|
| 74 |
|
| 75 |
// Preprompt is the same as the first system message.
|
|
|
|
| 26 |
HarmBlockThreshold.BLOCK_ONLY_HIGH,
|
| 27 |
])
|
| 28 |
.optional(),
|
| 29 |
+
tools: z.array(z.any()),
|
| 30 |
});
|
| 31 |
|
| 32 |
export function endpointVertex(input: z.input<typeof endpointVertexParametersSchema>): Endpoint {
|
| 33 |
+
const { project, location, model, apiEndpoint, safetyThreshold, tools } =
|
| 34 |
endpointVertexParametersSchema.parse(input);
|
| 35 |
|
| 36 |
const vertex_ai = new VertexAI({
|
|
|
|
| 40 |
});
|
| 41 |
|
| 42 |
return async ({ messages, preprompt, generateSettings }) => {
|
| 43 |
+
const parameters = { ...model.parameters, ...generateSettings };
|
| 44 |
+
|
| 45 |
const generativeModel = vertex_ai.getGenerativeModel({
|
| 46 |
model: model.id ?? model.name,
|
| 47 |
safetySettings: safetyThreshold
|
|
|
|
| 69 |
]
|
| 70 |
: undefined,
|
| 71 |
generationConfig: {
|
| 72 |
+
maxOutputTokens: parameters?.max_new_tokens ?? 4096,
|
| 73 |
+
stopSequences: parameters?.stop,
|
| 74 |
+
temperature: parameters?.temperature ?? 1,
|
| 75 |
},
|
| 76 |
+
tools,
|
| 77 |
});
|
| 78 |
|
| 79 |
// Preprompt is the same as the first system message.
|