diff --git a/packages/google-vertex/src/google-vertex-language-model.test.ts b/packages/google-vertex/src/google-vertex-language-model.test.ts index bb8780bae78..268a001810b 100644 --- a/packages/google-vertex/src/google-vertex-language-model.test.ts +++ b/packages/google-vertex/src/google-vertex-language-model.test.ts @@ -153,6 +153,7 @@ describe('doGenerate', () => { frequencyPenalty: undefined, maxOutputTokens: undefined, responseMimeType: undefined, + responseSchema: undefined, temperature: undefined, topK: undefined, topP: undefined, @@ -234,6 +235,7 @@ describe('doGenerate', () => { generationConfig: { maxOutputTokens: 100, responseMimeType: undefined, + responseSchema: undefined, temperature: 0.5, topK: 0.1, topP: 0.9, @@ -267,6 +269,7 @@ describe('doGenerate', () => { frequencyPenalty: undefined, maxOutputTokens: undefined, responseMimeType: undefined, + responseSchema: undefined, stopSequences: undefined, temperature: undefined, topK: undefined, @@ -321,46 +324,102 @@ describe('doGenerate', () => { }); }); - it('should set name & description in object-json mode', async () => { + it('should pass specification in object-json mode with structuredOutputs = true (default)', async () => { const { model, mockVertexAI } = createModel({ - modelId: 'test-model', generateContent: prepareResponse({ - parts: [{ text: '{"value":"Spark"}' }], + text: '{"property1":"value1","property2":"value2"}', }), }); - const response = await model.doGenerate({ + const result = await model.doGenerate({ inputFormat: 'prompt', mode: { type: 'object-json', - name: 'test-name', - description: 'test description', schema: { type: 'object', - properties: { value: { type: 'string' } }, - required: ['value'], + properties: { + property1: { type: 'string' }, + property2: { type: 'number' }, + }, + required: ['property1', 'property2'], + additionalProperties: false, + }, + }, + prompt: TEST_PROMPT, + }); + + expect(mockVertexAI.lastModelParams).toStrictEqual({ + generationConfig: { + frequencyPenalty: undefined, + maxOutputTokens: undefined, + responseMimeType: 'application/json', + responseSchema: { + properties: { + property1: { type: 'string' }, + property2: { type: 'number' }, + }, + required: ['property1', 'property2'], + type: 'object', + }, + stopSequences: undefined, + temperature: undefined, + topK: undefined, + topP: undefined, + }, + model: 'gemini-1.0-pro-002', + safetySettings: undefined, + }); + + expect(result.text).toStrictEqual( + '{"property1":"value1","property2":"value2"}', + ); + }); + + it('should not pass specification in object-json mode with structuredOutputs = false', async () => { + const { model, mockVertexAI } = createModel({ + generateContent: prepareResponse({ + text: '{"property1":"value1","property2":"value2"}', + }), + settings: { + structuredOutputs: false, + }, + }); + + const result = await model.doGenerate({ + inputFormat: 'prompt', + mode: { + type: 'object-json', + schema: { + type: 'object', + properties: { + property1: { type: 'string' }, + property2: { type: 'number' }, + }, + required: ['property1', 'property2'], additionalProperties: false, - $schema: 'http://json-schema.org/draft-07/schema#', }, }, prompt: TEST_PROMPT, }); expect(mockVertexAI.lastModelParams).toStrictEqual({ - model: 'test-model', generationConfig: { frequencyPenalty: undefined, maxOutputTokens: undefined, responseMimeType: 'application/json', + responseSchema: undefined, stopSequences: undefined, temperature: undefined, topK: undefined, topP: undefined, }, + model: 'gemini-1.0-pro-002', safetySettings: undefined, }); - expect(response.text).toStrictEqual('{"value":"Spark"}'); + expect(result.text).toStrictEqual( + '{"property1":"value1","property2":"value2"}', + ); }); it('should support object-tool mode', async () => { @@ -403,6 +462,7 @@ describe('doGenerate', () => { frequencyPenalty: undefined, maxOutputTokens: undefined, responseMimeType: undefined, + responseSchema: undefined, temperature: undefined, topK: undefined, topP: undefined, diff --git a/packages/google-vertex/src/google-vertex-language-model.ts b/packages/google-vertex/src/google-vertex-language-model.ts index 2879d34e8b5..d9a267d0f5d 100644 --- a/packages/google-vertex/src/google-vertex-language-model.ts +++ b/packages/google-vertex/src/google-vertex-language-model.ts @@ -13,6 +13,7 @@ import { GenerateContentResponse, GenerationConfig, Part, + ResponseSchema, SafetySetting, Tool, ToolConfig, @@ -33,9 +34,13 @@ type GoogleVertexAIConfig = { export class GoogleVertexLanguageModel implements LanguageModelV1 { readonly specificationVersion = 'v1'; readonly provider = 'google-vertex'; - readonly defaultObjectGenerationMode = 'tool'; + readonly defaultObjectGenerationMode = 'json'; readonly supportsImageUrls = false; + get supportsObjectGeneration() { + return this.settings.structuredOutputs !== false; + } + readonly modelId: GoogleVertexModelId; readonly settings: GoogleVertexSettings; @@ -98,8 +103,20 @@ export class GoogleVertexLanguageModel implements LanguageModelV1 { temperature, topP, stopSequences, + + // response format: responseMimeType: responseFormat?.type === 'json' ? 'application/json' : undefined, + responseSchema: + responseFormat?.type === 'json' && + responseFormat.schema != null && + // Google Vertex does not support all OpenAPI Schema features, + // so this is needed as an escape hatch: + this.supportsObjectGeneration + ? (convertJSONSchemaToOpenAPISchema( + responseFormat.schema, + ) as ResponseSchema) + : undefined, }; const type = mode.type; @@ -132,6 +149,15 @@ export class GoogleVertexLanguageModel implements LanguageModelV1 { generationConfig: { ...generationConfig, responseMimeType: 'application/json', + responseSchema: + mode.schema != null && + // Google Vertex does not support all OpenAPI Schema features, + // so this is needed as an escape hatch: + this.supportsObjectGeneration + ? (convertJSONSchemaToOpenAPISchema( + mode.schema, + ) as ResponseSchema) + : undefined, }, safetySettings: this.settings.safetySettings as | undefined diff --git a/packages/google-vertex/src/google-vertex-settings.ts b/packages/google-vertex/src/google-vertex-settings.ts index 06b0382cd1f..6e76ca5eb3a 100644 --- a/packages/google-vertex/src/google-vertex-settings.ts +++ b/packages/google-vertex/src/google-vertex-settings.ts @@ -18,6 +18,16 @@ Models running with nucleus sampling don't allow topK setting. */ topK?: number; + /** + * Optional. Enable structured output. Default is true. + * + * This is useful when the JSON Schema contains elements that are + * not supported by the OpenAPI schema version that + * Google Generative AI uses. You can use this to disable + * structured outputs if you need to. + */ + structuredOutputs?: boolean; + /** Optional. A list of unique safety settings for blocking unsafe content. */