Skip to content

Commit ded26b7

Browse files
hsubox76hsubox76
committed
fill out Schema builder
1 parent 2a16ff5 commit ded26b7

File tree

5 files changed

+234
-107
lines changed

5 files changed

+234
-107
lines changed

‎common/api-review/vertexai-preview.api.md

Lines changed: 6 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -203,42 +203,8 @@ export interface FunctionCallPart {
203203
export interface FunctionDeclaration {
204204
description?: string;
205205
name: string;
206-
parameters?: FunctionDeclarationSchema;
207-
}
208-
209-
// @public
210-
export interface FunctionDeclarationSchema {
211-
description?: string;
212-
properties: {
213-
[k: string]: FunctionDeclarationSchemaProperty;
214-
};
215-
required?: string[];
216-
type: FunctionDeclarationSchemaType;
217-
}
218-
219-
// @public
220-
export interface FunctionDeclarationSchemaProperty {
221-
description?: string;
222-
enum?: string[];
223-
example?: unknown;
224-
format?: string;
225-
items?: FunctionDeclarationSchema;
226-
nullable?: boolean;
227-
properties?: {
228-
[k: string]: FunctionDeclarationSchema;
229-
};
230-
required?: string[];
231-
type?: FunctionDeclarationSchemaType;
232-
}
233-
234-
// @public
235-
export enum FunctionDeclarationSchemaType {
236-
ARRAY = "ARRAY",
237-
BOOLEAN = "BOOLEAN",
238-
INTEGER = "INTEGER",
239-
NUMBER = "NUMBER",
240-
OBJECT = "OBJECT",
241-
STRING = "STRING"
206+
// Warning: (ae-forgotten-export) The symbol "ObjectSchema" needs to be exported by the entry point index.d.ts
207+
parameters?: ObjectSchema;
242208
}
243209

244210
// @public
@@ -331,6 +297,10 @@ export interface GenerationConfig {
331297
// (undocumented)
332298
presencePenalty?: number;
333299
responseMimeType?: string;
300+
// Warning: (ae-forgotten-export) The symbol "TypedSchema" needs to be exported by the entry point index.d.ts
301+
//
302+
// (undocumented)
303+
responseSchema?: TypedSchema;
334304
// (undocumented)
335305
stopSequences?: string[];
336306
// (undocumented)

‎packages/vertexai/src/api.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import { VertexAIError } from './errors';
2626
import { GenerativeModel } from './models/generative-model';
2727

2828
export { ChatSession } from './methods/chat-session';
29+
export * from './requests/schema-builder';
2930

3031
export { GenerativeModel };
3132

‎packages/vertexai/src/requests/schema-builder.test.ts

Lines changed: 171 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,28 +17,189 @@
1717

1818
import { expect, use } from 'chai';
1919
import sinonChai from 'sinon-chai';
20-
import { ObjectSchema, Schema } from './schema-builder';
20+
import { Schema } from './schema-builder';
2121

2222
use(sinonChai);
2323

24-
describe.only('request formatting methods', () => {
24+
describe.only('Schema builder', () => {
25+
it('builds integer schema', () => {
26+
const schema = Schema.createInteger();
27+
expect(schema.toJSON()).to.eql({
28+
type: 'integer',
29+
nullable: false
30+
});
31+
});
32+
it('builds number schema', () => {
33+
const schema = Schema.createNumber();
34+
expect(schema.toJSON()).to.eql({
35+
type: 'number',
36+
nullable: false
37+
});
38+
});
39+
it('builds boolean schema', () => {
40+
const schema = Schema.createBoolean();
41+
expect(schema.toJSON()).to.eql({
42+
type: 'boolean',
43+
nullable: false
44+
});
45+
});
2546
it('builds string schema', () => {
2647
const schema = Schema.createString({ description: 'hey' });
27-
expect(schema.toJSON().type).to.equal('string');
28-
expect(schema.toJSON().description).to.equal('hey');
48+
expect(schema.toJSON()).to.eql({
49+
type: 'string',
50+
description: 'hey',
51+
nullable: false
52+
});
53+
});
54+
it('builds enumString schema', () => {
55+
const schema = Schema.createEnumString({
56+
example: 'east',
57+
enumValues: ['east', 'west']
58+
});
59+
expect(schema.toJSON()).to.eql({
60+
type: 'string',
61+
example: 'east',
62+
enum: ['east', 'west'],
63+
nullable: false
64+
});
65+
});
66+
it('builds functionDeclaration schema', () => {
67+
const schema = Schema.createFunctionDeclaration({
68+
properties: {
69+
'someInput': Schema.createString()
70+
}
71+
});
72+
expect(schema.toJSON()).to.eql({
73+
type: 'object',
74+
nullable: false,
75+
properties: {
76+
'someInput': {
77+
type: 'string',
78+
nullable: false
79+
}
80+
},
81+
required: ['someInput']
82+
});
2983
});
3084
it('builds layered schema', () => {
3185
const schema = Schema.createArray({
3286
items: Schema.createObject({
3387
properties: {
34-
name: Schema.createString({})
88+
country: Schema.createString({
89+
description: 'some country',
90+
required: true
91+
}),
92+
population: Schema.createInteger(),
93+
coordinates: Schema.createObject({
94+
properties: {
95+
latitude: Schema.createNumber({ format: 'float' }),
96+
longitude: Schema.createNumber({ format: 'double' })
97+
}
98+
}),
99+
hemisphere: Schema.createObject({
100+
properties: {
101+
latitudinal: Schema.createEnumString({ enumValues: ['N', 'S'] }),
102+
longitudinal: Schema.createEnumString({ enumValues: ['E', 'W'] })
103+
}
104+
}),
105+
isCapital: Schema.createBoolean()
35106
}
36107
})
37108
});
38-
expect(schema.toJSON().type).to.equal('array');
39-
expect((schema.toJSON().items as Schema).type).to.equal('object');
40-
expect(
41-
((schema.toJSON().items as ObjectSchema).properties.name as Schema).type
42-
).to.equal('string');
109+
110+
expect(schema.toJSON()).to.eql(layeredSchemaOutput);
111+
});
112+
it('can override the "required" and "nullable" properties', () => {
113+
const schema = Schema.createObject({
114+
properties: {
115+
country: Schema.createString(),
116+
elevation: Schema.createNumber({ required: false }),
117+
population: Schema.createInteger({ nullable: true })
118+
}
119+
});
120+
expect(schema.toJSON()).to.eql({
121+
'type': 'object',
122+
'nullable': false,
123+
'properties': {
124+
'country': {
125+
'type': 'string',
126+
'nullable': false
127+
},
128+
'elevation': {
129+
'type': 'number',
130+
'nullable': false
131+
},
132+
'population': {
133+
'type': 'integer',
134+
'nullable': true
135+
}
136+
},
137+
'required': ['country', 'population']
138+
});
43139
});
44140
});
141+
142+
const layeredSchemaOutput = {
143+
'type': 'array',
144+
'nullable': false,
145+
'items': {
146+
'type': 'object',
147+
'nullable': false,
148+
'properties': {
149+
'country': {
150+
'type': 'string',
151+
'description': 'some country',
152+
'nullable': false
153+
},
154+
'population': {
155+
'type': 'integer',
156+
'nullable': false
157+
},
158+
'coordinates': {
159+
'type': 'object',
160+
'nullable': false,
161+
'properties': {
162+
'latitude': {
163+
'type': 'number',
164+
'format': 'float',
165+
'nullable': false
166+
},
167+
'longitude': {
168+
'type': 'number',
169+
'format': 'double',
170+
'nullable': false
171+
}
172+
},
173+
'required': ['latitude', 'longitude']
174+
},
175+
'hemisphere': {
176+
'type': 'object',
177+
'nullable': false,
178+
'properties': {
179+
'latitudinal': {
180+
'type': 'string',
181+
'nullable': false,
182+
'enum': ['N', 'S']
183+
},
184+
'longitudinal': {
185+
'type': 'string',
186+
'nullable': false,
187+
'enum': ['E', 'W']
188+
}
189+
},
190+
'required': ['latitudinal', 'longitudinal']
191+
},
192+
'isCapital': {
193+
'type': 'boolean',
194+
'nullable': false
195+
}
196+
},
197+
'required': [
198+
'country',
199+
'population',
200+
'coordinates',
201+
'hemisphere',
202+
'isCapital'
203+
]
204+
}
205+
};

‎packages/vertexai/src/requests/schema-builder.ts

Lines changed: 53 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { SchemaInterface, SchemaType, SchemaParams } from '../types/schema';
22

3-
export class Schema implements SchemaInterface {
3+
export abstract class Schema implements SchemaInterface {
44
/**
55
* Optional. The type of the property. {@link
66
* SchemaType}.
@@ -15,9 +15,9 @@ export class Schema implements SchemaInterface {
1515
format?: string;
1616
/** Optional. The description of the property. */
1717
description?: string;
18-
/** Optional. Whether the property is nullable. */
18+
/** Optional. Whether the property is nullable. Defaults to false. */
1919
nullable: boolean;
20-
/** Optional. Array of required property. */
20+
/** Optional. Defaults to true. */
2121
required: boolean;
2222
/** Optional. The example of the property. */
2323
example?: unknown;
@@ -27,15 +27,22 @@ export class Schema implements SchemaInterface {
2727
this.format = schemaParams?.format;
2828
this.description = schemaParams?.description;
2929
this.example = schemaParams?.example;
30-
this.nullable = schemaParams?.nullable || false;
31-
this.required = schemaParams?.required || true;
30+
this.nullable = schemaParams.hasOwnProperty('nullable')
31+
? !!schemaParams.nullable
32+
: false;
33+
this.required = schemaParams.hasOwnProperty('required')
34+
? !!schemaParams.required
35+
: true;
3236
}
3337

38+
/** Converts class to a plain JSON object (not a string). */
3439
toJSON(): Record<string, unknown> {
3540
const obj: Record<string, unknown> = {};
3641
for (const prop in this) {
3742
if (this.hasOwnProperty(prop) && this[prop] !== undefined) {
38-
obj[prop] = this[prop];
43+
if (prop !== 'required' || this.type === SchemaType.OBJECT) {
44+
obj[prop] = this[prop];
45+
}
3946
}
4047
}
4148
return obj;
@@ -57,7 +64,17 @@ export class Schema implements SchemaInterface {
5764
return new ObjectSchema(objectParams, objectParams.properties);
5865
}
5966

60-
static createString(stringParams: SchemaParams): StringSchema {
67+
static createFunctionDeclaration(
68+
objectParams: SchemaParams & {
69+
properties: {
70+
[k: string]: Schema;
71+
};
72+
}
73+
): ObjectSchema {
74+
return this.createObject(objectParams);
75+
}
76+
77+
static createString(stringParams?: SchemaParams): StringSchema {
6178
return new StringSchema(stringParams);
6279
}
6380

@@ -80,6 +97,14 @@ export class Schema implements SchemaInterface {
8097
}
8198
}
8299

100+
export type TypedSchema =
101+
| IntegerSchema
102+
| NumberSchema
103+
| StringSchema
104+
| BooleanSchema
105+
| ObjectSchema
106+
| ArraySchema;
107+
83108
export class IntegerSchema extends Schema {
84109
constructor(schemaParams?: SchemaParams) {
85110
super({
@@ -112,10 +137,19 @@ export class StringSchema extends Schema {
112137
...schemaParams
113138
});
114139
}
140+
141+
toJSON(): Record<string, unknown> {
142+
const obj = super.toJSON();
143+
if (this.enumValues) {
144+
obj['enum'] = this.enumValues;
145+
delete obj.enumValues;
146+
}
147+
return obj;
148+
}
115149
}
116150

117151
export class ArraySchema extends Schema {
118-
constructor(schemaParams: SchemaParams, public items: Schema) {
152+
constructor(schemaParams: SchemaParams, public items: TypedSchema) {
119153
super({
120154
type: SchemaType.ARRAY,
121155
...schemaParams
@@ -133,7 +167,7 @@ export class ObjectSchema extends Schema {
133167
constructor(
134168
schemaParams: SchemaParams,
135169
public properties: {
136-
[k: string]: Schema;
170+
[k: string]: TypedSchema;
137171
}
138172
) {
139173
super({
@@ -145,12 +179,19 @@ export class ObjectSchema extends Schema {
145179
toJSON(): Record<string, unknown> {
146180
const obj = super.toJSON();
147181
const properties: Record<string, unknown> = {};
148-
for (const property in this.properties) {
149-
if (this.properties.hasOwnProperty(property)) {
150-
properties[property] = this.properties[property].toJSON();
182+
const required = [];
183+
for (const propertyKey in this.properties) {
184+
if (this.properties.hasOwnProperty(propertyKey)) {
185+
properties[propertyKey] = this.properties[propertyKey].toJSON();
186+
if (this.properties[propertyKey].required) {
187+
required.push(propertyKey);
188+
}
151189
}
152190
}
153191
obj.properties = properties;
192+
if (required.length > 0) {
193+
obj.required = required;
194+
}
154195
return obj;
155196
}
156197
}

0 commit comments

Comments
 (0)