Repository navigation
[inference] Sanitize tool schemas for Google Vertex AI inference endpoints #284810
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| import { pick } from 'lodash'; | ||
| import { z } from '@kbn/zod/v4'; | ||
| import type { ToolSchema } from '@kbn/inference-common'; | ||
| import { sanitizeToolSchemasForVertex } from './sanitize_tool_schemas_for_vertex'; | ||
|
|
||
| // mirrors how tool schemas are converted before reaching the adapters, | ||
| // see `resolveToolSchema` in `@kbn/inference-langchain` | ||
| const toToolSchema = (zodSchema: z.ZodType): ToolSchema => | ||
| pick(z.toJSONSchema(zodSchema, { io: 'input' }), [ | ||
| 'type', | ||
| 'properties', | ||
| 'required', | ||
| ]) as ToolSchema; | ||
|
|
||
| describe('sanitizeToolSchemasForVertex', () => { | ||
| it('rebuilds tool schemas keeping only the fields Vertex AI accepts', () => { | ||
| const schema = toToolSchema( | ||
| z.object({ | ||
| size: z.number().positive().lt(100).optional(), | ||
| mode: z.literal('fast'), | ||
| filter: z.union([z.string(), z.array(z.string())]).optional(), | ||
| note: z.string().nullable(), | ||
| data: z.record(z.string(), z.string()), | ||
| // discriminated unions emit `oneOf` with `const` discriminators inside | ||
| operations: z.array( | ||
| z.discriminatedUnion('type', [ | ||
| z.object({ type: z.literal('add'), config: z.record(z.string(), z.string()) }), | ||
| z.object({ type: z.literal('remove') }), | ||
| ]) | ||
| ), | ||
| }) | ||
| ); | ||
| // zod v4 emits keywords rejected by Vertex: `propertyNames` for records, | ||
| // `exclusiveMinimum`/`exclusiveMaximum` for number bounds and `const` for literals | ||
| expect(JSON.stringify(schema)).toContain('propertyNames'); | ||
| expect(JSON.stringify(schema)).toContain('exclusiveMinimum'); | ||
| expect(JSON.stringify(schema)).toContain('exclusiveMaximum'); | ||
| expect(JSON.stringify(schema)).toContain('const'); | ||
|
|
||
| const sanitized = sanitizeToolSchemasForVertex({ | ||
| myTool: { | ||
| description: 'some cool tool', | ||
| schema, | ||
| }, | ||
| }); | ||
|
|
||
| expect(sanitized).toEqual({ | ||
| myTool: { | ||
| description: 'some cool tool', | ||
| schema: { | ||
| type: 'object', | ||
| properties: { | ||
| size: { | ||
| type: 'number', | ||
| }, | ||
| mode: { | ||
| type: 'string', | ||
| enum: ['fast'], | ||
| }, | ||
| filter: { | ||
| anyOf: [{ type: 'string' }, { type: 'array', items: { type: 'string' } }], | ||
| }, | ||
| note: { | ||
| type: 'string', | ||
| nullable: true, | ||
| }, | ||
| data: { | ||
| type: 'object', | ||
| additionalProperties: { type: 'string' }, | ||
| }, | ||
| operations: { | ||
| type: 'array', | ||
| items: { | ||
| anyOf: [ | ||
| { | ||
| type: 'object', | ||
| properties: { | ||
| type: { type: 'string', enum: ['add'] }, | ||
| config: { | ||
| type: 'object', | ||
| additionalProperties: { type: 'string' }, | ||
| }, | ||
| }, | ||
| required: ['type', 'config'], | ||
| }, | ||
| { | ||
| type: 'object', | ||
| properties: { | ||
| type: { type: 'string', enum: ['remove'] }, | ||
| }, | ||
| required: ['type'], | ||
| }, | ||
| ], | ||
| }, | ||
| }, | ||
| }, | ||
| required: ['mode', 'note', 'data', 'operations'], | ||
| }, | ||
| }, | ||
| }); | ||
| }); | ||
| }); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,123 @@ | ||
| /* | ||
|
Contributor
Author
There was a problem hiding this comment. Comments in this file are verbose, but considering what it does, I think it's worth keeping them to be explicitly clear about why these functions are needed. Will remove if anyone thinks otherwise. |
||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| import { mapValues } from 'lodash'; | ||
| import type { ToolOptions, ToolSchemaType } from '@kbn/inference-common'; | ||
|
|
||
| export const sanitizeToolSchemasForVertex = (tools: ToolOptions['tools']): ToolOptions['tools'] => { | ||
| if (!tools) { | ||
| return tools; | ||
| } | ||
|
|
||
| return mapValues(tools, (tool) => ({ | ||
| ...tool, | ||
| schema: tool.schema ? toVertexSchema(tool.schema) : undefined, | ||
| })); | ||
| }; | ||
|
|
||
| /** | ||
| * Fields of Vertex AI's function declaration schema, "a select subset of an | ||
| * OpenAPI 3.0 schema object". Any other field is rejected with a 400. | ||
| * https://github.com/googleapis/googleapis/blob/master/google/cloud/aiplatform/v1/openapi.proto | ||
| */ | ||
| const VERTEX_SCHEMA_FIELDS = new Set([ | ||
| 'type', | ||
| 'format', | ||
| 'title', | ||
| 'description', | ||
| 'nullable', | ||
| 'default', | ||
| 'items', | ||
| 'minItems', | ||
| 'maxItems', | ||
| 'enum', | ||
| 'properties', | ||
| 'propertyOrdering', | ||
| 'required', | ||
| 'minProperties', | ||
| 'maxProperties', | ||
| 'minimum', | ||
| 'maximum', | ||
| 'minLength', | ||
| 'maxLength', | ||
| 'pattern', | ||
| 'example', | ||
| 'anyOf', | ||
| 'additionalProperties', | ||
| ]); | ||
|
|
||
| /** | ||
| * `ToolSchemaType` does not model everything zod v4 emits, so the | ||
| * conversion operates on this loose representation. | ||
| */ | ||
| interface SchemaPart { | ||
| [key: string]: unknown; | ||
| const?: unknown; | ||
| type?: string; | ||
| nullable?: boolean; | ||
| enum?: unknown[]; | ||
| properties?: Record<string, ToolSchemaType>; | ||
| items?: ToolSchemaType; | ||
| additionalProperties?: ToolSchemaType | boolean; | ||
| anyOf?: ToolSchemaType[]; | ||
| oneOf?: ToolSchemaType[]; | ||
| } | ||
|
|
||
| /** | ||
| * Rebuilds a tool schema keeping only the fields Vertex AI accepts, so that | ||
| * JSON Schema keywords zod v4 (or user-provided tool schemas) emit never reach | ||
| * Vertex's parser. Conversions preserve the intent of dropped fields: | ||
| * `const: 'x'` becomes `enum: ['x']`, `oneOf` branches become `anyOf`, and a | ||
| * `null` branch in `anyOf` (zod's representation of nullable values) becomes | ||
| * `nullable: true`. | ||
| */ | ||
| const toVertexSchema = <T extends ToolSchemaType>(schemaPart: T): T => { | ||
|
Comment on lines
+70
to
+78
Contributor
There was a problem hiding this comment. Yeah I really hate that we have to do this, it's crazy that they didn't improve this in 2 years, but I don't see an alternative, so... |
||
| const source = schemaPart as SchemaPart; | ||
| const schema: SchemaPart = {}; | ||
|
|
||
| for (const [key, value] of Object.entries(source)) { | ||
| if (VERTEX_SCHEMA_FIELDS.has(key)) { | ||
| schema[key] = value; | ||
| } | ||
| } | ||
|
|
||
| if (typeof source.const === 'string' && schema.enum === undefined) { | ||
| schema.enum = [source.const]; | ||
| } | ||
|
|
||
| if (source.oneOf) { | ||
| schema.anyOf = [...(schema.anyOf ?? []), ...source.oneOf]; | ||
| } | ||
|
|
||
| if (schema.properties) { | ||
| schema.properties = mapValues(schema.properties, toVertexSchema); | ||
| } | ||
| if (schema.items) { | ||
| schema.items = toVertexSchema(schema.items); | ||
| } | ||
| if (typeof schema.additionalProperties === 'object') { | ||
| schema.additionalProperties = toVertexSchema(schema.additionalProperties); | ||
| } | ||
| if (schema.anyOf) { | ||
| const branches = schema.anyOf.map(toVertexSchema); | ||
| const nonNullBranches = branches.filter((branch) => (branch as SchemaPart).type !== 'null'); | ||
|
|
||
| if (nonNullBranches.length === branches.length) { | ||
| schema.anyOf = branches; | ||
| } else { | ||
| schema.nullable = true; | ||
| if (nonNullBranches.length === 1) { | ||
| delete schema.anyOf; | ||
| Object.assign(schema, nonNullBranches[0]); | ||
| } else { | ||
| schema.anyOf = nonNullBranches; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return schema as unknown as T; | ||
| }; | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would move that inside
createEndpointRequest, better isolation of concerns ihmoThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Simple condition, if it ever expands to more providers, definitely makes sense, gonna merge now to unblock.