Add some test cases, fix error

Author: hsubox76

packages/vertexai/src/errors.ts

@@ -33,8 +33,8 @@ export class VertexAIError extends FirebaseError {
    * @param customErrorData - Optional error data.
    */
   constructor(
-    readonly code: VertexAIErrorCode,
-    readonly message: string,
+    code: VertexAIErrorCode,
+    message: string,
     readonly customErrorData?: CustomErrorData
   ) {
     // Match error format used by FirebaseError from ErrorFactory

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

@@ -18,6 +18,7 @@
 import { expect, use } from 'chai';
 import sinonChai from 'sinon-chai';
 import { Schema } from './schema-builder';
+import { VertexAIErrorCode } from '../types';
 
 use(sinonChai);
 
@@ -29,13 +30,30 @@ describe.only('Schema builder', () => {
       nullable: false
     });
   });
+  it('builds integer schema with options and overrides', () => {
+    const schema = Schema.integer({ nullable: true, format: 'int32' });
+    expect(schema.toRequest()).to.eql({
+      type: 'integer',
+      format: 'int32',
+      nullable: true
+    });
+  });
   it('builds number schema', () => {
     const schema = Schema.number();
     expect(schema.toRequest()).to.eql({
       type: 'number',
       nullable: false
     });
   });
+  it('builds number schema with options and unknown options', () => {
+    const schema = Schema.number({ format: 'float', futureOption: 'test' });
+    expect(schema.toRequest()).to.eql({
+      type: 'number',
+      format: 'float',
+      futureOption: 'test',
+      nullable: false
+    });
+  });
   it('builds boolean schema', () => {
     const schema = Schema.boolean();
     expect(schema.toRequest()).to.eql({
@@ -81,6 +99,48 @@ describe.only('Schema builder', () => {
       required: ['someInput']
     });
   });
+  it('builds an object schema', () => {
+    const schema = Schema.object({
+      properties: {
+        'someInput': Schema.string()
+      }
+    });
+    expect(schema.toRequest()).to.eql({
+      type: 'object',
+      nullable: false,
+      properties: {
+        'someInput': {
+          type: 'string',
+          nullable: false
+        }
+      },
+      required: ['someInput']
+    });
+  });
+  it('builds an object schema with optional properties', () => {
+    const schema = Schema.object({
+      properties: {
+        'someInput': Schema.string(),
+        'someBool': Schema.boolean()
+      },
+      optionalProperties: ['someBool']
+    });
+    expect(schema.toRequest()).to.eql({
+      type: 'object',
+      nullable: false,
+      properties: {
+        'someInput': {
+          type: 'string',
+          nullable: false
+        },
+        'someBool': {
+          type: 'boolean',
+          nullable: false
+        }
+      },
+      required: ['someInput']
+    });
+  });
   it('builds layered schema - partially filled out', () => {
     const schema = Schema.array({
       items: Schema.object({
@@ -192,6 +252,17 @@ describe.only('Schema builder', () => {
       'required': ['country', 'population']
     });
   });
+  it('throws if an optionalProperties item does not exist', () => {
+    const schema = Schema.object({
+      properties: {
+        country: Schema.string(),
+        elevation: Schema.number(),
+        population: Schema.integer({ nullable: true })
+      },
+      optionalProperties: ['cat']
+    });
+    expect(() => schema.toRequest()).to.throw(VertexAIErrorCode.INVALID_SCHEMA);
+  });
 });
 
 const layeredSchemaOutputPartial = {

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

@@ -1,3 +1,5 @@
+import { VertexAIError } from '../errors';
+import { VertexAIErrorCode } from '../types';
 import {
   SchemaInterface,
   SchemaType,
@@ -25,12 +27,19 @@ export abstract class Schema implements SchemaInterface {
   nullable: boolean;
   /** Optional. The example of the property. */
   example?: unknown;
+  /**
+   * Allows user to add other schema properties that have not yet
+   * been officially added to the SDK.
+   */
+  [key: string]: unknown;
 
   constructor(schemaParams: SchemaInterface) {
+    // eslint-disable-next-line guard-for-in
+    for (const paramKey in schemaParams) {
+      this[paramKey] = schemaParams[paramKey];
+    }
+    // Ensure these are explicitly set to avoid TS errors.
     this.type = schemaParams.type;
-    this.format = schemaParams?.format;
-    this.description = schemaParams?.description;
-    this.example = schemaParams?.example;
     this.nullable = schemaParams.hasOwnProperty('nullable')
       ? !!schemaParams.nullable
       : false;
@@ -127,6 +136,7 @@ export class IntegerSchema extends Schema {
     });
   }
 }
+
 export class NumberSchema extends Schema {
   constructor(schemaParams?: SchemaParams) {
     super({
@@ -135,6 +145,7 @@ export class NumberSchema extends Schema {
     });
   }
 }
+
 export class BooleanSchema extends Schema {
   constructor(schemaParams?: SchemaParams) {
     super({
@@ -194,19 +205,28 @@ export class ObjectSchema extends Schema {
 
   toRequest(): _SchemaRequest {
     const obj = super.toRequest();
-    const properties: Record<string, _SchemaRequest> = {};
+    obj.properties = this.properties;
     const required = [];
+    if (this.optionalProperties) {
+      for (const propertyKey of this.optionalProperties) {
+        if (!this.properties.hasOwnProperty(propertyKey)) {
+          throw new VertexAIError(
+            VertexAIErrorCode.INVALID_SCHEMA,
+            `Property "${propertyKey}" specified in "optionalProperties" does not exist.`
+          );
+        }
+      }
+    }
     for (const propertyKey in this.properties) {
       if (this.properties.hasOwnProperty(propertyKey)) {
-        properties[propertyKey] = this.properties[
+        obj.properties[propertyKey] = this.properties[
           propertyKey
         ].toRequest() as _SchemaRequest;
         if (!this.optionalProperties.includes(propertyKey)) {
           required.push(propertyKey);
         }
       }
     }
-    obj.properties = properties;
     if (required.length > 0) {
       obj.required = required;
     }

packages/vertexai/src/types/error.ts

@@ -78,6 +78,9 @@ export const enum VertexAIErrorCode {
   /** An error associated with a Content object.  */
   INVALID_CONTENT = 'invalid-content',
 
+  /** An error due to invalid Schema input.  */
+  INVALID_SCHEMA = 'invalid-schema',
+
   /** An error occurred due to a missing Firebase API key. */
   NO_API_KEY = 'no-api-key',
 

packages/vertexai/src/types/schema.ts

@@ -35,6 +35,7 @@ interface SchemaShared<T> {
   example?: unknown;
   /** Optional. Whether the property is nullable. */
   nullable?: boolean;
+  [key: string]: unknown;
 }
 
 /**