Skip to content

Commit 85ff2c4

Browse files
lahirumarambahsubox76
lahirumaramba
and
hsubox76
authoredMay 29, 2024··
Add responseSchema to generationConfig (#158)
* Add JSON Schema Support * Add unit tests * Add changeset * update changeset * Add some unit test cases (#159) --------- Co-authored-by: Christina Holland <chholland@google.com>
·
v0.24.0@google/generative-ai@0.12.0
1 parent bb7e6c9 commit 85ff2c4

23 files changed

+261
-101
lines changed
 

‎.changeset/five-spoons-matter.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@google/generative-ai": minor
3+
---
4+
5+
Added `responseSchema` to `GenerationConfig` to allow user to provide a JSON schema when `responseMimeType` is set to JSON.

‎docs/contributing.md

+1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ Steps to build locally:
3737
1. Run `yarn test` to run unit tests.
3838
1. Run `yarn docs` to generate any changes to reference docs (destination dir is docs/reference).
3939
1. Run `yarn format` to fix formatting and add license headers as needed.
40+
1. Run `yarn changeset` to generate the changeset summary.
4041

4142
This repo has a monorepo structure to allow for easily adding additional packages. The `@google/generative-ai` package code is in `packages/main`.
4243

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.description.md

-13
This file was deleted.

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.example.md

-13
This file was deleted.

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.format.md

-13
This file was deleted.

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.md

+3-16
Original file line numberDiff line numberDiff line change
@@ -4,25 +4,12 @@
44

55
## FunctionDeclarationSchemaProperty interface
66

7-
Schema is used to define the format of input/output data. Represents a select subset of an OpenAPI 3.0 schema object. More fields may be added in the future as needed.
7+
Schema for top-level function declaration
88

99
**Signature:**
1010

1111
```typescript
12-
export interface FunctionDeclarationSchemaProperty
12+
export interface FunctionDeclarationSchemaProperty extends Schema
1313
```
14-
15-
## Properties
16-
17-
| Property | Modifiers | Type | Description |
18-
| --- | --- | --- | --- |
19-
| [description?](./generative-ai.functiondeclarationschemaproperty.description.md) | | string | _(Optional)_ Optional. The description of the property. |
20-
| [enum?](./generative-ai.functiondeclarationschemaproperty.enum.md) | | string\[\] | _(Optional)_ Optional. The enum of the property. |
21-
| [example?](./generative-ai.functiondeclarationschemaproperty.example.md) | | unknown | _(Optional)_ Optional. The example of the property. |
22-
| [format?](./generative-ai.functiondeclarationschemaproperty.format.md) | | string | _(Optional)_ Optional. The format of the property. |
23-
| [items?](./generative-ai.functiondeclarationschemaproperty.items.md) | | [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md) | _(Optional)_ Optional. The items of the property. [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md) |
24-
| [nullable?](./generative-ai.functiondeclarationschemaproperty.nullable.md) | | boolean | _(Optional)_ Optional. Whether the property is nullable. |
25-
| [properties?](./generative-ai.functiondeclarationschemaproperty.properties.md) | | { \[k: string\]: [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md)<!-- -->; } | _(Optional)_ Optional. Map of [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md)<!-- -->. |
26-
| [required?](./generative-ai.functiondeclarationschemaproperty.required.md) | | string\[\] | _(Optional)_ Optional. Array of required property. |
27-
| [type?](./generative-ai.functiondeclarationschemaproperty.type.md) | | [FunctionDeclarationSchemaType](./generative-ai.functiondeclarationschematype.md) | _(Optional)_ Optional. The type of the property. [FunctionDeclarationSchemaType](./generative-ai.functiondeclarationschematype.md)<!-- -->. |
14+
**Extends:** [Schema](./generative-ai.schema.md)
2815

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.required.md

-13
This file was deleted.

‎docs/reference/main/generative-ai.generationconfig.md

+1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export interface GenerationConfig
1919
| [candidateCount?](./generative-ai.generationconfig.candidatecount.md) | | number | _(Optional)_ |
2020
| [maxOutputTokens?](./generative-ai.generationconfig.maxoutputtokens.md) | | number | _(Optional)_ |
2121
| [responseMimeType?](./generative-ai.generationconfig.responsemimetype.md) | | string | _(Optional)_ Output response mimetype of the generated candidate text. Supported mimetype: <code>text/plain</code>: (default) Text output. <code>application/json</code>: JSON response in the candidates. |
22+
| [responseSchema?](./generative-ai.generationconfig.responseschema.md) | | [ResponseSchema](./generative-ai.responseschema.md) | _(Optional)_ Output response schema of the generated candidate text. Note: This only applies when the specified <code>responseMIMEType</code> supports a schema; currently this is limited to <code>application/json</code>. |
2223
| [stopSequences?](./generative-ai.generationconfig.stopsequences.md) | | string\[\] | _(Optional)_ |
2324
| [temperature?](./generative-ai.generationconfig.temperature.md) | | number | _(Optional)_ |
2425
| [topK?](./generative-ai.generationconfig.topk.md) | | number | _(Optional)_ |

‎docs/reference/main/generative-ai.generationconfig.responseschema.md

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [GenerationConfig](./generative-ai.generationconfig.md) &gt; [responseSchema](./generative-ai.generationconfig.responseschema.md)
4+
5+
## GenerationConfig.responseSchema property
6+
7+
Output response schema of the generated candidate text. Note: This only applies when the specified `responseMIMEType` supports a schema; currently this is limited to `application/json`<!-- -->.
8+
9+
**Signature:**
10+
11+
```typescript
12+
responseSchema?: ResponseSchema;
13+
```

‎docs/reference/main/generative-ai.md

+3-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@
5353
| [FunctionCallPart](./generative-ai.functioncallpart.md) | Content part interface if the part represents FunctionResponse. |
5454
| [FunctionDeclaration](./generative-ai.functiondeclaration.md) | Structured representation of a function declaration as defined by the \[OpenAPI 3.0 specification\](https://spec.openapis.org/oas/v3.0.3). Included in this declaration are the function name and parameters. This FunctionDeclaration is a representation of a block of code that can be used as a Tool by the model and executed by the client. |
5555
| [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md) | Schema for parameters passed to [FunctionDeclaration.parameters](./generative-ai.functiondeclaration.parameters.md)<!-- -->. |
56-
| [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) | Schema is used to define the format of input/output data. Represents a select subset of an OpenAPI 3.0 schema object. More fields may be added in the future as needed. |
56+
| [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) | Schema for top-level function declaration |
5757
| [FunctionDeclarationsTool](./generative-ai.functiondeclarationstool.md) | A FunctionDeclarationsTool is a piece of code that enables the system to interact with external systems to perform an action, or set of actions, outside of knowledge and scope of the model. |
5858
| [FunctionResponse](./generative-ai.functionresponse.md) | The result output from a \[FunctionCall\] that contains a string representing the \[FunctionDeclaration.name\] and a structured JSON object containing any output from the function is used as context to the model. This should contain the result of a \[FunctionCall\] made based on model prediction. |
5959
| [FunctionResponsePart](./generative-ai.functionresponsepart.md) | Content part interface if the part represents FunctionResponse. |
@@ -68,8 +68,10 @@
6868
| [ModelParams](./generative-ai.modelparams.md) | Params passed to [GoogleGenerativeAI.getGenerativeModel()](./generative-ai.googlegenerativeai.getgenerativemodel.md)<!-- -->. |
6969
| [PromptFeedback](./generative-ai.promptfeedback.md) | If the prompt was blocked, this will be populated with <code>blockReason</code> and the relevant <code>safetyRatings</code>. |
7070
| [RequestOptions](./generative-ai.requestoptions.md) | Params passed to getGenerativeModel() or GoogleAIFileManager(). |
71+
| [ResponseSchema](./generative-ai.responseschema.md) | Schema passed to [GenerationConfig.responseSchema](./generative-ai.generationconfig.responseschema.md) |
7172
| [SafetyRating](./generative-ai.safetyrating.md) | A safety rating associated with a [GenerateContentCandidate](./generative-ai.generatecontentcandidate.md) |
7273
| [SafetySetting](./generative-ai.safetysetting.md) | Safety setting that can be sent as part of request parameters. |
74+
| [Schema](./generative-ai.schema.md) | Schema is used to define the format of input/output data. Represents a select subset of an OpenAPI 3.0 schema object. More fields may be added in the future as needed. |
7375
| [StartChatParams](./generative-ai.startchatparams.md) | Params for [GenerativeModel.startChat()](./generative-ai.generativemodel.startchat.md)<!-- -->. |
7476
| [TextPart](./generative-ai.textpart.md) | Content part interface if the part represents a text string. |
7577
| [ToolConfig](./generative-ai.toolconfig.md) | Tool config. This config is shared for all tools provided in the request. |

‎docs/reference/main/generative-ai.responseschema.md

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [ResponseSchema](./generative-ai.responseschema.md)
4+
5+
## ResponseSchema interface
6+
7+
Schema passed to [GenerationConfig.responseSchema](./generative-ai.generationconfig.responseschema.md)
8+
9+
**Signature:**
10+
11+
```typescript
12+
export interface ResponseSchema extends Schema
13+
```
14+
**Extends:** [Schema](./generative-ai.schema.md)
15+

‎docs/reference/main/generative-ai.schema.description.md

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [description](./generative-ai.schema.description.md)
4+
5+
## Schema.description property
6+
7+
Optional. The description of the property.
8+
9+
**Signature:**
10+
11+
```typescript
12+
description?: string;
13+
```

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.enum.md ‎docs/reference/main/generative-ai.schema.enum.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
22

3-
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) &gt; [enum](./generative-ai.functiondeclarationschemaproperty.enum.md)
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [enum](./generative-ai.schema.enum.md)
44

5-
## FunctionDeclarationSchemaProperty.enum property
5+
## Schema.enum property
66

77
Optional. The enum of the property.
88

‎docs/reference/main/generative-ai.schema.example.md

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [example](./generative-ai.schema.example.md)
4+
5+
## Schema.example property
6+
7+
Optional. The example of the property.
8+
9+
**Signature:**
10+
11+
```typescript
12+
example?: unknown;
13+
```

‎docs/reference/main/generative-ai.schema.format.md

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [format](./generative-ai.schema.format.md)
4+
5+
## Schema.format property
6+
7+
Optional. The format of the property.
8+
9+
**Signature:**
10+
11+
```typescript
12+
format?: string;
13+
```

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.items.md ‎docs/reference/main/generative-ai.schema.items.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
22

3-
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) &gt; [items](./generative-ai.functiondeclarationschemaproperty.items.md)
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [items](./generative-ai.schema.items.md)
44

5-
## FunctionDeclarationSchemaProperty.items property
5+
## Schema.items property
66

77
Optional. The items of the property. [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md)
88

‎docs/reference/main/generative-ai.schema.md

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md)
4+
5+
## Schema interface
6+
7+
Schema is used to define the format of input/output data. Represents a select subset of an OpenAPI 3.0 schema object. More fields may be added in the future as needed.
8+
9+
**Signature:**
10+
11+
```typescript
12+
export interface Schema
13+
```
14+
15+
## Properties
16+
17+
| Property | Modifiers | Type | Description |
18+
| --- | --- | --- | --- |
19+
| [description?](./generative-ai.schema.description.md) | | string | _(Optional)_ Optional. The description of the property. |
20+
| [enum?](./generative-ai.schema.enum.md) | | string\[\] | _(Optional)_ Optional. The enum of the property. |
21+
| [example?](./generative-ai.schema.example.md) | | unknown | _(Optional)_ Optional. The example of the property. |
22+
| [format?](./generative-ai.schema.format.md) | | string | _(Optional)_ Optional. The format of the property. |
23+
| [items?](./generative-ai.schema.items.md) | | [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md) | _(Optional)_ Optional. The items of the property. [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md) |
24+
| [nullable?](./generative-ai.schema.nullable.md) | | boolean | _(Optional)_ Optional. Whether the property is nullable. |
25+
| [properties?](./generative-ai.schema.properties.md) | | { \[k: string\]: [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md)<!-- -->; } | _(Optional)_ Optional. Map of [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md)<!-- -->. |
26+
| [required?](./generative-ai.schema.required.md) | | string\[\] | _(Optional)_ Optional. Array of required property. |
27+
| [type?](./generative-ai.schema.type.md) | | [FunctionDeclarationSchemaType](./generative-ai.functiondeclarationschematype.md) | _(Optional)_ Optional. The type of the property. [FunctionDeclarationSchemaType](./generative-ai.functiondeclarationschematype.md)<!-- -->. |
28+

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.nullable.md ‎docs/reference/main/generative-ai.schema.nullable.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
22

3-
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) &gt; [nullable](./generative-ai.functiondeclarationschemaproperty.nullable.md)
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [nullable](./generative-ai.schema.nullable.md)
44

5-
## FunctionDeclarationSchemaProperty.nullable property
5+
## Schema.nullable property
66

77
Optional. Whether the property is nullable.
88

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.properties.md ‎docs/reference/main/generative-ai.schema.properties.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
22

3-
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) &gt; [properties](./generative-ai.functiondeclarationschemaproperty.properties.md)
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [properties](./generative-ai.schema.properties.md)
44

5-
## FunctionDeclarationSchemaProperty.properties property
5+
## Schema.properties property
66

77
Optional. Map of [FunctionDeclarationSchema](./generative-ai.functiondeclarationschema.md)<!-- -->.
88

‎docs/reference/main/generative-ai.schema.required.md

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
2+
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [required](./generative-ai.schema.required.md)
4+
5+
## Schema.required property
6+
7+
Optional. Array of required property.
8+
9+
**Signature:**
10+
11+
```typescript
12+
required?: string[];
13+
```

‎docs/reference/main/generative-ai.functiondeclarationschemaproperty.type.md ‎docs/reference/main/generative-ai.schema.type.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
<!-- Do not edit this file. It is automatically generated by API Documenter. -->
22

3-
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [FunctionDeclarationSchemaProperty](./generative-ai.functiondeclarationschemaproperty.md) &gt; [type](./generative-ai.functiondeclarationschemaproperty.type.md)
3+
[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [Schema](./generative-ai.schema.md) &gt; [type](./generative-ai.schema.type.md)
44

5-
## FunctionDeclarationSchemaProperty.type property
5+
## Schema.type property
66

77
Optional. The type of the property. [FunctionDeclarationSchemaType](./generative-ai.functiondeclarationschematype.md)<!-- -->.
88

‎packages/main/src/models/generative-model.test.ts

+96-6
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import { GenerativeModel } from "./generative-model";
1919
import * as sinonChai from "sinon-chai";
2020
import {
2121
FunctionCallingMode,
22+
FunctionDeclarationSchemaType,
2223
HarmBlockThreshold,
2324
HarmCategory,
2425
} from "../../types";
@@ -50,7 +51,19 @@ describe("GenerativeModel", () => {
5051
"apiKey",
5152
{
5253
model: "my-model",
53-
generationConfig: { temperature: 0 },
54+
generationConfig: {
55+
temperature: 0,
56+
responseMimeType: "application/json",
57+
responseSchema: {
58+
type: FunctionDeclarationSchemaType.OBJECT,
59+
properties: {
60+
testField: {
61+
type: FunctionDeclarationSchemaType.STRING,
62+
properties: {},
63+
},
64+
},
65+
},
66+
},
5467
safetySettings: [
5568
{
5669
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
@@ -68,6 +81,15 @@ describe("GenerativeModel", () => {
6881
},
6982
);
7083
expect(genModel.generationConfig?.temperature).to.equal(0);
84+
expect(genModel.generationConfig?.responseMimeType).to.equal(
85+
"application/json",
86+
);
87+
expect(genModel.generationConfig?.responseSchema.type).to.equal(
88+
FunctionDeclarationSchemaType.OBJECT,
89+
);
90+
expect(
91+
genModel.generationConfig?.responseSchema.properties.testField.type,
92+
).to.equal(FunctionDeclarationSchemaType.STRING);
7193
expect(genModel.safetySettings?.length).to.equal(1);
7294
expect(genModel.tools?.length).to.equal(1);
7395
expect(genModel.toolConfig?.functionCallingConfig.mode).to.equal(
@@ -92,6 +114,7 @@ describe("GenerativeModel", () => {
92114
value.includes(FunctionCallingMode.NONE) &&
93115
value.includes("be friendly") &&
94116
value.includes("temperature") &&
117+
value.includes("testField") &&
95118
value.includes(HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
96119
);
97120
}),
@@ -129,7 +152,19 @@ describe("GenerativeModel", () => {
129152
it("generateContent overrides model values", async () => {
130153
const genModel = new GenerativeModel("apiKey", {
131154
model: "my-model",
132-
generationConfig: { temperature: 0 },
155+
generationConfig: {
156+
temperature: 0,
157+
responseMimeType: "application/json",
158+
responseSchema: {
159+
type: FunctionDeclarationSchemaType.OBJECT,
160+
properties: {
161+
testField: {
162+
type: FunctionDeclarationSchemaType.STRING,
163+
properties: {},
164+
},
165+
},
166+
},
167+
},
133168
safetySettings: [
134169
{
135170
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
@@ -152,7 +187,18 @@ describe("GenerativeModel", () => {
152187
mockResponse as Response,
153188
);
154189
await genModel.generateContent({
155-
generationConfig: { topK: 1 },
190+
generationConfig: {
191+
topK: 1,
192+
responseSchema: {
193+
type: FunctionDeclarationSchemaType.OBJECT,
194+
properties: {
195+
newTestField: {
196+
type: FunctionDeclarationSchemaType.STRING,
197+
properties: {},
198+
},
199+
},
200+
},
201+
},
156202
safetySettings: [
157203
{
158204
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
@@ -179,6 +225,8 @@ describe("GenerativeModel", () => {
179225
value.includes(FunctionCallingMode.AUTO) &&
180226
value.includes("be formal") &&
181227
value.includes("topK") &&
228+
value.includes("newTestField") &&
229+
!value.includes("testField") &&
182230
value.includes(HarmCategory.HARM_CATEGORY_HARASSMENT)
183231
);
184232
}),
@@ -203,7 +251,6 @@ describe("GenerativeModel", () => {
203251
mockResponse as Response,
204252
);
205253
await genModel.countTokens("hello");
206-
console.log(makeRequestStub.args[0]);
207254
expect(makeRequestStub).to.be.calledWith(
208255
"models/my-model",
209256
request.Task.COUNT_TOKENS,
@@ -254,9 +301,24 @@ describe("GenerativeModel", () => {
254301
it("passes params through to chat.sendMessage", async () => {
255302
const genModel = new GenerativeModel("apiKey", {
256303
model: "my-model",
304+
generationConfig: {
305+
temperature: 0,
306+
responseMimeType: "application/json",
307+
responseSchema: {
308+
type: FunctionDeclarationSchemaType.OBJECT,
309+
properties: {
310+
testField: {
311+
type: FunctionDeclarationSchemaType.STRING,
312+
properties: {},
313+
},
314+
},
315+
},
316+
},
257317
systemInstruction: { role: "system", parts: [{ text: "be friendly" }] },
258318
});
259319
expect(genModel.systemInstruction?.parts[0].text).to.equal("be friendly");
320+
expect(genModel.generationConfig.responseSchema.properties.testField).to
321+
.exist;
260322
const mockResponse = getMockResponse(
261323
"unary-success-basic-reply-short.json",
262324
);
@@ -270,7 +332,7 @@ describe("GenerativeModel", () => {
270332
match.any,
271333
false,
272334
match((value: string) => {
273-
return value.includes("be friendly");
335+
return value.includes("be friendly") && value.includes("testField");
274336
}),
275337
{},
276338
);
@@ -279,10 +341,25 @@ describe("GenerativeModel", () => {
279341
it("startChat overrides model values", async () => {
280342
const genModel = new GenerativeModel("apiKey", {
281343
model: "my-model",
344+
generationConfig: {
345+
temperature: 0,
346+
responseMimeType: "application/json",
347+
responseSchema: {
348+
type: FunctionDeclarationSchemaType.OBJECT,
349+
properties: {
350+
testField: {
351+
type: FunctionDeclarationSchemaType.STRING,
352+
properties: {},
353+
},
354+
},
355+
},
356+
},
282357
tools: [{ functionDeclarations: [{ name: "myfunc" }] }],
283358
toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
284359
systemInstruction: { role: "system", parts: [{ text: "be friendly" }] },
285360
});
361+
expect(genModel.generationConfig.responseSchema.properties.testField).to
362+
.exist;
286363
expect(genModel.tools?.length).to.equal(1);
287364
expect(genModel.toolConfig?.functionCallingConfig.mode).to.equal(
288365
FunctionCallingMode.NONE,
@@ -297,6 +374,17 @@ describe("GenerativeModel", () => {
297374
await genModel
298375
.startChat({
299376
tools: [{ functionDeclarations: [{ name: "otherfunc" }] }],
377+
generationConfig: {
378+
responseSchema: {
379+
type: FunctionDeclarationSchemaType.OBJECT,
380+
properties: {
381+
newTestField: {
382+
type: FunctionDeclarationSchemaType.STRING,
383+
properties: {},
384+
},
385+
},
386+
},
387+
},
300388
toolConfig: {
301389
functionCallingConfig: { mode: FunctionCallingMode.AUTO },
302390
},
@@ -312,7 +400,9 @@ describe("GenerativeModel", () => {
312400
return (
313401
value.includes("otherfunc") &&
314402
value.includes(FunctionCallingMode.AUTO) &&
315-
value.includes("be formal")
403+
value.includes("be formal") &&
404+
value.includes("newTestField") &&
405+
!value.includes("testField")
316406
);
317407
}),
318408
{},

‎packages/main/types/requests.ts

+34-16
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,12 @@ export interface GenerationConfig {
8181
* `application/json`: JSON response in the candidates.
8282
*/
8383
responseMimeType?: string;
84+
/**
85+
* Output response schema of the generated candidate text.
86+
* Note: This only applies when the specified `responseMIMEType` supports a schema; currently
87+
* this is limited to `application/json`.
88+
*/
89+
responseSchema?: ResponseSchema;
8490
}
8591

8692
/**
@@ -238,28 +244,13 @@ export enum FunctionDeclarationSchemaType {
238244
OBJECT = "OBJECT",
239245
}
240246

241-
/**
242-
* Schema for parameters passed to {@link FunctionDeclaration.parameters}.
243-
* @public
244-
*/
245-
export interface FunctionDeclarationSchema {
246-
/** The type of the parameter. */
247-
type: FunctionDeclarationSchemaType;
248-
/** The format of the parameter. */
249-
properties: { [k: string]: FunctionDeclarationSchemaProperty };
250-
/** Optional. Description of the parameter. */
251-
description?: string;
252-
/** Optional. Array of required parameters. */
253-
required?: string[];
254-
}
255-
256247
/**
257248
* Schema is used to define the format of input/output data.
258249
* Represents a select subset of an OpenAPI 3.0 schema object.
259250
* More fields may be added in the future as needed.
260251
* @public
261252
*/
262-
export interface FunctionDeclarationSchemaProperty {
253+
export interface Schema {
263254
/**
264255
* Optional. The type of the property. {@link
265256
* FunctionDeclarationSchemaType}.
@@ -283,6 +274,33 @@ export interface FunctionDeclarationSchemaProperty {
283274
example?: unknown;
284275
}
285276

277+
/**
278+
* Schema for parameters passed to {@link FunctionDeclaration.parameters}.
279+
* @public
280+
*/
281+
export interface FunctionDeclarationSchema {
282+
/** The type of the parameter. */
283+
type: FunctionDeclarationSchemaType;
284+
/** The format of the parameter. */
285+
properties: { [k: string]: FunctionDeclarationSchemaProperty };
286+
/** Optional. Description of the parameter. */
287+
description?: string;
288+
/** Optional. Array of required parameters. */
289+
required?: string[];
290+
}
291+
292+
/**
293+
* Schema for top-level function declaration
294+
* @public
295+
*/
296+
export interface FunctionDeclarationSchemaProperty extends Schema {}
297+
298+
/**
299+
* Schema passed to {@link GenerationConfig.responseSchema}
300+
* @public
301+
*/
302+
export interface ResponseSchema extends Schema {}
303+
286304
/**
287305
* Tool config. This config is shared for all tools provided in the request.
288306
* @public

0 commit comments

Comments
 (0)
Please sign in to comment.