start work on Schema rewrite

hsubox76 · hsubox76 · commit 2a16ff50e7e1 · 2024-09-03T15:57:51.000-07:00
diff --git a/packages/vertexai/src/requests/schema-builder.test.ts b/packages/vertexai/src/requests/schema-builder.test.ts
@@ -0,0 +1,44 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { expect, use } from 'chai';
+import sinonChai from 'sinon-chai';
+import { ObjectSchema, Schema } from './schema-builder';
+
+use(sinonChai);
+
+describe.only('request formatting methods', () => {
+  it('builds string schema', () => {
+    const schema = Schema.createString({ description: 'hey' });
+    expect(schema.toJSON().type).to.equal('string');
+    expect(schema.toJSON().description).to.equal('hey');
+  });
+  it('builds layered schema', () => {
+    const schema = Schema.createArray({
+      items: Schema.createObject({
+        properties: {
+          name: Schema.createString({})
+        }
+      })
+    });
+    expect(schema.toJSON().type).to.equal('array');
+    expect((schema.toJSON().items as Schema).type).to.equal('object');
+    expect(
+      ((schema.toJSON().items as ObjectSchema).properties.name as Schema).type
+    ).to.equal('string');
+  });
+});
diff --git a/packages/vertexai/src/requests/schema-builder.ts b/packages/vertexai/src/requests/schema-builder.ts
@@ -0,0 +1,156 @@
+import { SchemaInterface, SchemaType, SchemaParams } from '../types/schema';
+
+export class Schema implements SchemaInterface {
+  /**
+   * Optional. The type of the property. {@link
+   * SchemaType}.
+   */
+  type: SchemaType;
+  /** Optional. The format of the property.
+   * Supported formats:
+   *  for NUMBER type: "float", "double"
+   *  for INTEGER type: "int32", "int64"
+   *  for STRING type: "email", "byte", etc
+   */
+  format?: string;
+  /** Optional. The description of the property. */
+  description?: string;
+  /** Optional. Whether the property is nullable. */
+  nullable: boolean;
+  /** Optional. Array of required property. */
+  required: boolean;
+  /** Optional. The example of the property. */
+  example?: unknown;
+
+  constructor(schemaParams: SchemaInterface) {
+    this.type = schemaParams.type;
+    this.format = schemaParams?.format;
+    this.description = schemaParams?.description;
+    this.example = schemaParams?.example;
+    this.nullable = schemaParams?.nullable || false;
+    this.required = schemaParams?.required || true;
+  }
+
+  toJSON(): Record<string, unknown> {
+    const obj: Record<string, unknown> = {};
+    for (const prop in this) {
+      if (this.hasOwnProperty(prop) && this[prop] !== undefined) {
+        obj[prop] = this[prop];
+      }
+    }
+    return obj;
+  }
+
+  static createArray(
+    arrayParams: SchemaParams & { items: Schema }
+  ): ArraySchema {
+    return new ArraySchema(arrayParams, arrayParams.items);
+  }
+
+  static createObject(
+    objectParams: SchemaParams & {
+      properties: {
+        [k: string]: Schema;
+      };
+    }
+  ): ObjectSchema {
+    return new ObjectSchema(objectParams, objectParams.properties);
+  }
+
+  static createString(stringParams: SchemaParams): StringSchema {
+    return new StringSchema(stringParams);
+  }
+
+  static createEnumString(
+    stringParams: SchemaParams & { enumValues: string[] }
+  ): StringSchema {
+    return new StringSchema(stringParams, stringParams.enumValues);
+  }
+
+  static createInteger(integerParams?: SchemaParams): IntegerSchema {
+    return new IntegerSchema(integerParams);
+  }
+
+  static createNumber(numberParams?: SchemaParams): NumberSchema {
+    return new NumberSchema(numberParams);
+  }
+
+  static createBoolean(booleanParams?: SchemaParams): BooleanSchema {
+    return new BooleanSchema(booleanParams);
+  }
+}
+
+export class IntegerSchema extends Schema {
+  constructor(schemaParams?: SchemaParams) {
+    super({
+      type: SchemaType.INTEGER,
+      ...schemaParams
+    });
+  }
+}
+export class NumberSchema extends Schema {
+  constructor(schemaParams?: SchemaParams) {
+    super({
+      type: SchemaType.NUMBER,
+      ...schemaParams
+    });
+  }
+}
+export class BooleanSchema extends Schema {
+  constructor(schemaParams?: SchemaParams) {
+    super({
+      type: SchemaType.BOOLEAN,
+      ...schemaParams
+    });
+  }
+}
+
+export class StringSchema extends Schema {
+  constructor(schemaParams?: SchemaParams, public enumValues?: string[]) {
+    super({
+      type: SchemaType.STRING,
+      ...schemaParams
+    });
+  }
+}
+
+export class ArraySchema extends Schema {
+  constructor(schemaParams: SchemaParams, public items: Schema) {
+    super({
+      type: SchemaType.ARRAY,
+      ...schemaParams
+    });
+  }
+
+  toJSON(): Record<string, unknown> {
+    const obj = super.toJSON();
+    obj.items = this.items.toJSON();
+    return obj;
+  }
+}
+
+export class ObjectSchema extends Schema {
+  constructor(
+    schemaParams: SchemaParams,
+    public properties: {
+      [k: string]: Schema;
+    }
+  ) {
+    super({
+      type: SchemaType.OBJECT,
+      ...schemaParams
+    });
+  }
+
+  toJSON(): Record<string, unknown> {
+    const obj = super.toJSON();
+    const properties: Record<string, unknown> = {};
+    for (const property in this.properties) {
+      if (this.properties.hasOwnProperty(property)) {
+        properties[property] = this.properties[property].toJSON();
+      }
+    }
+    obj.properties = properties;
+    return obj;
+  }
+}
diff --git a/packages/vertexai/src/types/requests.ts b/packages/vertexai/src/types/requests.ts
@@ -22,6 +22,7 @@ import {
   HarmBlockThreshold,
   HarmCategory
 } from './enums';
+import { SchemaType } from './schema';
 
 /**
  * Base parameters for a number of methods.
@@ -177,33 +178,15 @@ export declare interface FunctionDeclarationsTool {
   functionDeclarations?: FunctionDeclaration[];
 }
 
-/**
- * Contains the list of OpenAPI data types
- * as defined by https://swagger.io/docs/specification/data-models/data-types/
- * @public
- */
-export enum FunctionDeclarationSchemaType {
-  /** String type. */
-  STRING = 'STRING',
-  /** Number type. */
-  NUMBER = 'NUMBER',
-  /** Integer type. */
-  INTEGER = 'INTEGER',
-  /** Boolean type. */
-  BOOLEAN = 'BOOLEAN',
-  /** Array type. */
-  ARRAY = 'ARRAY',
-  /** Object type. */
-  OBJECT = 'OBJECT'
-}
-
 /**
  * Schema for parameters passed to {@link FunctionDeclaration.parameters}.
+ * For function declarations, the top level {@link Schema} type must be
+ * type `object`.
  * @public
  */
 export interface FunctionDeclarationSchema {
   /** The type of the parameter. */
-  type: FunctionDeclarationSchemaType;
+  type: SchemaType.OBJECT;
   /** The format of the parameter. */
   properties: { [k: string]: FunctionDeclarationSchemaProperty };
   /** Optional. Description of the parameter. */
@@ -221,9 +204,9 @@ export interface FunctionDeclarationSchema {
 export interface FunctionDeclarationSchemaProperty {
   /**
    * Optional. The type of the property. {@link
-   * FunctionDeclarationSchemaType}.
+   * SchemaType}.
    */
-  type?: FunctionDeclarationSchemaType;
+  type?: SchemaType;
   /** Optional. The format of the property. */
   format?: string;
   /** Optional. The description of the property. */
diff --git a/packages/vertexai/src/types/schema.ts b/packages/vertexai/src/types/schema.ts
@@ -0,0 +1,72 @@
+
+/**
+ * Contains the list of OpenAPI data types
+ * as defined by https://swagger.io/docs/specification/data-models/data-types/
+ * @public
+ */
+export enum SchemaType {
+  /** String type. */
+  STRING = "string",
+  /** Number type. */
+  NUMBER = "number",
+  /** Integer type. */
+  INTEGER = "integer",
+  /** Boolean type. */
+  BOOLEAN = "boolean",
+  /** Array type. */
+  ARRAY = "array",
+  /** Object type. */
+  OBJECT = "object",
+}
+
+interface SchemaShared {
+  /** Optional. The format of the property. */
+  format?: string;
+  /** Optional. The description of the property. */
+  description?: string;
+  /** Optional. The items of the property. */
+  items?: SchemaInterface;
+  /** Optional. The enum of the property. */
+  enum?: string[];
+  /** Optional. Map of {@link Schema}. */
+  properties?: {
+    [k: string]: SchemaInterface;
+  };
+  /** Optional. The example of the property. */
+  example?: unknown;
+  /** Optional. Whether the property is nullable. */
+  nullable?: boolean;
+}
+
+/**
+ * User-facing params passed to specific Schema static methods.
+ */
+export interface SchemaParams extends SchemaShared {
+  /** Optional. Array of required property. */
+  required?: boolean;
+
+}
+
+/**
+ * Final format for Schema params passed to backend requests.
+ */
+export interface SchemaRequest extends SchemaShared {
+  /**
+   * The type of the property. {@link
+   * SchemaType}.
+   */
+  type?: SchemaType;
+  /** Optional. Array of required property. */
+  required?: string[];
+}
+
+/**
+ * Interface for Schema class.
+ */
+export interface SchemaInterface extends SchemaParams {
+  /**
+   * The type of the property. {@link
+   * SchemaType}.
+   */
+  type: SchemaType;
+}
