Make Schema a discriminated union

This leverages the type system to better describe the API's requirements for schemas. For example, rather than saying that any schema might have an optional `items` property, we're able to express that `items` is required on array schemas and forbidden on all others. More info on discriminated unions: https://www.typescriptlang.org/docs/handbook/2/narrowing.html#discriminated-unions
google-gemini · Nov 13, 2024 · 0c417fb · 0c417fb
1 parent 6ec2c27
commit 0c417fb
Show file tree

Hide file tree

Showing 5 changed files with 260 additions and 65 deletions.
diff --git a/.changeset/young-rivers-shout.md b/.changeset/young-rivers-shout.md
@@ -0,0 +1,5 @@
+---
+"@google/generative-ai": minor
+---
+
+The schema types are now more specific, using a [discriminated union](https://www.typescriptlang.org/docs/handbook/2/narrowing.html#discriminated-unions) based on the 'type' field to more accurately define which fields are allowed.
diff --git a/common/api-review/generative-ai-server.api.md b/common/api-review/generative-ai-server.api.md
@@ -4,6 +4,33 @@
 
 ```ts
 
+// Warning: (ae-incompatible-release-tags) The symbol "ArraySchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface ArraySchema extends BaseSchema {
+    items: Schema;
+    maxItems?: number;
+    minItems?: number;
+    // (undocumented)
+    type: typeof SchemaType.ARRAY;
+}
+
+// Warning: (ae-internal-missing-underscore) The name "BaseSchema" should be prefixed with an underscore because the declaration is marked as @internal
+//
+// @internal
+export interface BaseSchema {
+    description?: string;
+    nullable?: boolean;
+}
+
+// Warning: (ae-incompatible-release-tags) The symbol "BooleanSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface BooleanSchema extends BaseSchema {
+    // (undocumented)
+    type: typeof SchemaType.BOOLEAN;
+}
+
 // @public
 export interface CachedContent extends CachedContentBase {
     createTime?: string;
@@ -286,8 +313,7 @@ export interface FunctionDeclarationSchema {
 }
 
 // @public
-export interface FunctionDeclarationSchemaProperty extends Schema {
-}
+export type FunctionDeclarationSchemaProperty = Schema;
 
 // @public
 export interface FunctionDeclarationsTool {
@@ -368,6 +394,15 @@ export interface InlineDataPart {
     text?: never;
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "IntegerSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface IntegerSchema extends BaseSchema {
+    format?: "int32" | "int64";
+    // (undocumented)
+    type: typeof SchemaType.INTEGER;
+}
+
 // @public (undocumented)
 export interface ListCacheResponse {
     // (undocumented)
@@ -392,6 +427,27 @@ export interface ListParams {
     pageToken?: string;
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "NumberSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface NumberSchema extends BaseSchema {
+    format?: "float" | "double";
+    // (undocumented)
+    type: typeof SchemaType.NUMBER;
+}
+
+// Warning: (ae-incompatible-release-tags) The symbol "ObjectSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface ObjectSchema extends BaseSchema {
+    properties: {
+        [k: string]: Schema;
+    };
+    required?: string[];
+    // (undocumented)
+    type: typeof SchemaType.OBJECT;
+}
+
 // @public
 export enum Outcome {
     OUTCOME_DEADLINE_EXCEEDED = "outcome_deadline_exceeded",
@@ -413,8 +469,7 @@ export interface RequestOptions {
 }
 
 // @public
-export interface ResponseSchema extends Schema {
-}
+export type ResponseSchema = Schema;
 
 // @public
 export interface RpcStatus {
@@ -424,19 +479,7 @@ export interface RpcStatus {
 }
 
 // @public
-export interface Schema {
-    description?: string;
-    enum?: string[];
-    example?: unknown;
-    format?: string;
-    items?: Schema;
-    nullable?: boolean;
-    properties?: {
-        [k: string]: Schema;
-    };
-    required?: string[];
-    type?: SchemaType;
-}
+export type Schema = StringSchema | NumberSchema | IntegerSchema | BooleanSchema | ArraySchema | ObjectSchema;
 
 // @public
 export enum SchemaType {
@@ -453,6 +496,15 @@ export interface SingleRequestOptions extends RequestOptions {
     signal?: AbortSignal;
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "StringSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface StringSchema extends BaseSchema {
+    enum?: string[];
+    // (undocumented)
+    type: typeof SchemaType.STRING;
+}
+
 // @public
 export interface TextPart {
     // (undocumented)

diff --git a/common/api-review/generative-ai.api.md b/common/api-review/generative-ai.api.md
@@ -4,6 +4,17 @@
 
 ```ts
 
+// Warning: (ae-incompatible-release-tags) The symbol "ArraySchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface ArraySchema extends BaseSchema {
+    items: Schema;
+    maxItems?: number;
+    minItems?: number;
+    // (undocumented)
+    type: typeof SchemaType.ARRAY;
+}
+
 // @public
 export interface BaseParams {
     // (undocumented)
@@ -12,6 +23,14 @@ export interface BaseParams {
     safetySettings?: SafetySetting[];
 }
 
+// Warning: (ae-internal-missing-underscore) The name "BaseSchema" should be prefixed with an underscore because the declaration is marked as @internal
+//
+// @internal
+export interface BaseSchema {
+    description?: string;
+    nullable?: boolean;
+}
+
 // @public
 export interface BatchEmbedContentsRequest {
     // (undocumented)
@@ -34,6 +53,14 @@ export enum BlockReason {
     SAFETY = "SAFETY"
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "BooleanSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface BooleanSchema extends BaseSchema {
+    // (undocumented)
+    type: typeof SchemaType.BOOLEAN;
+}
+
 // @public
 export interface CachedContent extends CachedContentBase {
     createTime?: string;
@@ -355,8 +382,7 @@ export interface FunctionDeclarationSchema {
 }
 
 // @public
-export interface FunctionDeclarationSchemaProperty extends Schema {
-}
+export type FunctionDeclarationSchemaProperty = Schema;
 
 // @public
 export interface FunctionDeclarationsTool {
@@ -645,6 +671,15 @@ export interface InlineDataPart {
     text?: never;
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "IntegerSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface IntegerSchema extends BaseSchema {
+    format?: "int32" | "int64";
+    // (undocumented)
+    type: typeof SchemaType.INTEGER;
+}
+
 // @public
 export interface LogprobsCandidate {
     logProbability: number;
@@ -672,6 +707,27 @@ export interface ModelParams extends BaseParams {
     tools?: Tool[];
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "NumberSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface NumberSchema extends BaseSchema {
+    format?: "float" | "double";
+    // (undocumented)
+    type: typeof SchemaType.NUMBER;
+}
+
+// Warning: (ae-incompatible-release-tags) The symbol "ObjectSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface ObjectSchema extends BaseSchema {
+    properties: {
+        [k: string]: Schema;
+    };
+    required?: string[];
+    // (undocumented)
+    type: typeof SchemaType.OBJECT;
+}
+
 // @public
 export enum Outcome {
     OUTCOME_DEADLINE_EXCEEDED = "outcome_deadline_exceeded",
@@ -706,8 +762,7 @@ export interface RequestOptions {
 }
 
 // @public
-export interface ResponseSchema extends Schema {
-}
+export type ResponseSchema = Schema;
 
 // @public
 export interface RetrievalMetadata {
@@ -731,19 +786,7 @@ export interface SafetySetting {
 }
 
 // @public
-export interface Schema {
-    description?: string;
-    enum?: string[];
-    example?: unknown;
-    format?: string;
-    items?: Schema;
-    nullable?: boolean;
-    properties?: {
-        [k: string]: Schema;
-    };
-    required?: string[];
-    type?: SchemaType;
-}
+export type Schema = StringSchema | NumberSchema | IntegerSchema | BooleanSchema | ArraySchema | ObjectSchema;
 
 // @public
 export enum SchemaType {
@@ -779,6 +822,15 @@ export interface StartChatParams extends BaseParams {
     tools?: Tool[];
 }
 
+// Warning: (ae-incompatible-release-tags) The symbol "StringSchema" is marked as @public, but its signature references "BaseSchema" which is marked as @internal
+//
+// @public
+export interface StringSchema extends BaseSchema {
+    enum?: string[];
+    // (undocumented)
+    type: typeof SchemaType.STRING;
+}
+
 // @public
 export enum TaskType {
     // (undocumented)

diff --git a/src/models/generative-model.test.ts b/src/models/generative-model.test.ts
@@ -22,6 +22,7 @@ import {
   FunctionCallingMode,
   HarmBlockThreshold,
   HarmCategory,
+  ObjectSchema,
   SchemaType,
 } from "../../types";
 import { getMockResponse } from "../../test-utils/mock-response";
@@ -60,7 +61,6 @@ describe("GenerativeModel", () => {
             properties: {
               testField: {
                 type: SchemaType.STRING,
-                properties: {},
               },
             },
           },
@@ -93,7 +93,8 @@ describe("GenerativeModel", () => {
       SchemaType.OBJECT,
     );
     expect(
-      genModel.generationConfig?.responseSchema.properties.testField.type,
+      (genModel.generationConfig?.responseSchema as ObjectSchema).properties
+        .testField.type,
     ).to.equal(SchemaType.STRING);
     expect(genModel.generationConfig?.presencePenalty).to.equal(0.6);
     expect(genModel.generationConfig?.frequencyPenalty).to.equal(0.5);
@@ -172,7 +173,6 @@ describe("GenerativeModel", () => {
           properties: {
             testField: {
               type: SchemaType.STRING,
-              properties: {},
             },
           },
         },
@@ -206,7 +206,6 @@ describe("GenerativeModel", () => {
           properties: {
             newTestField: {
               type: SchemaType.STRING,
-              properties: {},
             },
           },
         },
@@ -332,16 +331,17 @@ describe("GenerativeModel", () => {
           properties: {
             testField: {
               type: SchemaType.STRING,
-              properties: {},
             },
           },
         },
       },
       systemInstruction: { role: "system", parts: [{ text: "be friendly" }] },
     });
     expect(genModel.systemInstruction?.parts[0].text).to.equal("be friendly");
-    expect(genModel.generationConfig.responseSchema.properties.testField).to
-      .exist;
+    expect(
+      (genModel.generationConfig.responseSchema as ObjectSchema).properties
+        .testField,
+    ).to.exist;
     const mockResponse = getMockResponse(
       "unary-success-basic-reply-short.json",
     );
@@ -372,7 +372,6 @@ describe("GenerativeModel", () => {
           properties: {
             testField: {
               type: SchemaType.STRING,
-              properties: {},
             },
           },
         },
@@ -381,8 +380,10 @@ describe("GenerativeModel", () => {
       toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
       systemInstruction: { role: "system", parts: [{ text: "be friendly" }] },
     });
-    expect(genModel.generationConfig.responseSchema.properties.testField).to
-      .exist;
+    expect(
+      (genModel.generationConfig.responseSchema as ObjectSchema).properties
+        .testField,
+    ).to.exist;
     expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig.mode).to.equal(
       FunctionCallingMode.NONE,
@@ -403,7 +404,6 @@ describe("GenerativeModel", () => {
             properties: {
               newTestField: {
                 type: SchemaType.STRING,
-                properties: {},
               },
             },
           },