feat(appsync-modelgen-plugin): add support for generation route definitions (#869)

atierian · web-flow · commit 8c4be9adf8b2 · 2024-09-04T12:55:18.000-04:00
diff --git a/packages/appsync-modelgen-plugin/schemas/introspection/1/ModelIntrospectionSchema.json b/packages/appsync-modelgen-plugin/schemas/introspection/1/ModelIntrospectionSchema.json
@@ -26,6 +26,9 @@
         },
         "inputs": {
             "$ref": "#/definitions/SchemaInputs"
+        },
+        "generations": {
+            "$ref": "#/definitions/SchemaGenerations"
         }
     },
     "required": [
@@ -532,6 +535,9 @@
             ],
             "additionalProperties": false,
             "description": "Input Definition"
+        },
+        "SchemaGenerations": {
+            "$ref": "#/definitions/SchemaQueries"
         }
     }
 }
diff --git a/packages/appsync-modelgen-plugin/src/__tests__/visitors/__snapshots__/appsync-model-introspection-visitor.test.ts.snap b/packages/appsync-modelgen-plugin/src/__tests__/visitors/__snapshots__/appsync-model-introspection-visitor.test.ts.snap
@@ -1976,6 +1976,75 @@ exports[`Custom queries/mutations/subscriptions & input type tests should genera
 }"
 `;
 
+exports[`Generation Route Introspection Visitor Metadata snapshot should generate correct model intropection file validated by JSON schema 1`] = `
+"{
+    \\"version\\": 1,
+    \\"models\\": {},
+    \\"enums\\": {},
+    \\"nonModels\\": {
+        \\"Recipe\\": {
+            \\"name\\": \\"Recipe\\",
+            \\"fields\\": {
+                \\"name\\": {
+                    \\"name\\": \\"name\\",
+                    \\"isArray\\": false,
+                    \\"type\\": \\"String\\",
+                    \\"isRequired\\": false,
+                    \\"attributes\\": []
+                },
+                \\"ingredients\\": {
+                    \\"name\\": \\"ingredients\\",
+                    \\"isArray\\": true,
+                    \\"type\\": \\"String\\",
+                    \\"isRequired\\": false,
+                    \\"attributes\\": [],
+                    \\"isArrayNullable\\": true
+                },
+                \\"instructions\\": {
+                    \\"name\\": \\"instructions\\",
+                    \\"isArray\\": false,
+                    \\"type\\": \\"String\\",
+                    \\"isRequired\\": false,
+                    \\"attributes\\": []
+                }
+            }
+        }
+    },
+    \\"generations\\": {
+        \\"generateRecipe\\": {
+            \\"name\\": \\"generateRecipe\\",
+            \\"isArray\\": false,
+            \\"type\\": {
+                \\"nonModel\\": \\"Recipe\\"
+            },
+            \\"isRequired\\": false,
+            \\"arguments\\": {
+                \\"description\\": {
+                    \\"name\\": \\"description\\",
+                    \\"isArray\\": false,
+                    \\"type\\": \\"String\\",
+                    \\"isRequired\\": false
+                }
+            }
+        },
+        \\"summarize\\": {
+            \\"name\\": \\"summarize\\",
+            \\"isArray\\": false,
+            \\"type\\": \\"String\\",
+            \\"isRequired\\": false,
+            \\"arguments\\": {
+                \\"text\\": {
+                    \\"name\\": \\"text\\",
+                    \\"isArray\\": false,
+                    \\"type\\": \\"String\\",
+                    \\"isRequired\\": false
+                }
+            }
+        }
+    }
+}"
+`;
+
 exports[`Model Introspection Visitor Metadata snapshot should generate correct model intropection file validated by JSON schema 1`] = `
 "{
     \\"version\\": 1,
diff --git a/packages/appsync-modelgen-plugin/src/__tests__/visitors/appsync-model-introspection-visitor.test.ts b/packages/appsync-modelgen-plugin/src/__tests__/visitors/appsync-model-introspection-visitor.test.ts
@@ -1,6 +1,6 @@
 import { buildSchema, GraphQLSchema, parse, visit } from 'graphql';
 import { METADATA_SCALAR_MAP } from '../../scalars';
-import { AppSyncDirectives, DefaultDirectives, V1Directives, DeprecatedDirective, Directive } from '@aws-amplify/graphql-directives';
+import { AppSyncDirectives, DefaultDirectives, V1Directives, DeprecatedDirective, Directive, V2Directives } from '@aws-amplify/graphql-directives';
 import { scalars } from '../../scalars/supported-scalars';
 import { AppSyncModelIntrospectionVisitor } from '../../visitors/appsync-model-introspection-visitor';
 
@@ -806,3 +806,45 @@ describe('custom references', () => {
       .toThrowError(`Error processing @belongsTo directive on SqlRelated.primary. @hasOne or @hasMany directive with references ["primaryId"] was not found in connected model SqlPrimary`);
   });
 });
+
+describe('Generation Route Introspection Visitor', () => {
+  const schema = /* GraphQL */ `
+    type Recipe {
+      name: String
+      ingredients: [String]
+      instructions: String
+    }
+
+    type Query {
+      generateRecipe(description: String): Recipe
+      @generation(aiModel: "anthropic.claude-3-haiku-20240307-v1:0", systemPrompt: "You are a recipe generator.")
+
+      summarize(text: String): String
+      @generation(aiModel: "anthropic.claude-3-haiku-20240307-v1:0", systemPrompt: "You are a text summarizer.")
+    }
+  `;
+
+  const generationDirective: Directive = {
+    name: 'generation',
+    definition: /* GraphQL */ `
+      directive @generation(
+        aiModel: String!
+        systemPrompt: String!
+        inferenceConfiguration: GenerationInferenceConfiguration
+      ) on FIELD_DEFINITION
+
+      input GenerationInferenceConfiguration {
+        maxTokens: Int
+        temperature: Float
+        topP: Float
+      }
+    `,
+    defaults: {},
+  }
+  const visitor: AppSyncModelIntrospectionVisitor = getVisitor(schema, {}, [...V2Directives, generationDirective]);
+  describe('Metadata snapshot', () => {
+    it('should generate correct model intropection file validated by JSON schema', () => {
+      expect(visitor.generate()).toMatchSnapshot();
+    });
+  });
+});
diff --git a/packages/appsync-modelgen-plugin/src/interfaces/introspection/model-schema.ts b/packages/appsync-modelgen-plugin/src/interfaces/introspection/model-schema.ts
@@ -19,6 +19,7 @@
   mutations?: SchemaMutations;
   subscriptions?: SchemaSubscriptions;
   inputs?: SchemaInputs;
+  generations?: SchemaGenerations;
 };
 /**
  * Top-level Entities on a Schema
@@ -30,6 +31,7 @@ export type SchemaQueries = Record<string, SchemaQuery>;
 export type SchemaMutations = Record<string, SchemaMutation>;
 export type SchemaSubscriptions = Record<string, SchemaSubscription>;
 export type SchemaInputs = Record<string, Input>;
+export type SchemaGenerations = SchemaQueries;
 
 export type SchemaModel = {
   name: string;
diff --git a/packages/appsync-modelgen-plugin/src/utils/constants.ts b/packages/appsync-modelgen-plugin/src/utils/constants.ts
@@ -8,6 +8,7 @@ export const TransformerV2DirectiveName = {
   INDEX: 'index',
   DEFAULT: 'default',
   SEARCHABLE: 'searchable',
+  GENERATION: 'generation',
 };
 export const DEFAULT_HASH_KEY_FIELD = 'id';
 export const DEFAULT_CREATED_TIME = 'createdAt';
diff --git a/packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts b/packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts
@@ -1,4 +1,4 @@
-import { CodeGenDirective, CodeGenField, CodeGenModel } from '../visitors/appsync-visitor';
+import { CodeGenDirective, CodeGenField, CodeGenModel, CodeGenQuery } from '../visitors/appsync-visitor';
 import { TransformerV2DirectiveName } from './constants';
 
 export function addFieldToModel(model: CodeGenModel, field: CodeGenField): void {
@@ -40,4 +40,8 @@ export function getModelPrimaryKeyComponentFields(model: CodeGenModel): CodeGenF
     };
   }
   return keyFields;
+}
+
+export function containsGenerationDirective(queryField: CodeGenQuery): boolean {
+  return queryField.directives.some((directive) => directive.name === TransformerV2DirectiveName.GENERATION);
 }
diff --git a/packages/appsync-modelgen-plugin/src/validate-cjs.js b/packages/appsync-modelgen-plugin/src/validate-cjs.js
diff --git a/packages/appsync-modelgen-plugin/src/visitors/appsync-model-introspection-visitor.ts b/packages/appsync-modelgen-plugin/src/visitors/appsync-model-introspection-visitor.ts
@@ -4,6 +4,7 @@ import { Argument, AssociationType, Field, Fields, FieldType, ModelAttribute, Mo
 import { METADATA_SCALAR_MAP } from "../scalars";
 import { CodeGenConnectionType } from "../utils/process-connections";
 import { RawAppSyncModelConfig, ParsedAppSyncModelConfig, AppSyncModelVisitor, CodeGenEnum, CodeGenField, CodeGenModel, CodeGenPrimaryKeyType, CodeGenQuery, CodeGenSubscription, CodeGenMutation, CodeGenInputObject } from "./appsync-visitor";
+import { containsGenerationDirective } from "../utils/fieldUtils";
 
 const validateModelIntrospectionSchema = require('../validate-cjs');
 
@@ -65,12 +66,22 @@ export class AppSyncModelIntrospectionVisitor<
     const queries = Object.values(this.queryMap).reduce((acc, queryObj: CodeGenQuery) => {
       // Skip the field if the field type is union/interface
       // TODO: Remove this skip once these types are supported for stakeholder usages
+      // Also skip if this query has a generation directive. These are handled separately and keyed under `generations`.
       const fieldType = this.getType(queryObj.type) as any;
-      if (this.isUnionFieldType(fieldType) || this.isInterfaceFieldType(fieldType)) {
+      if (this.isUnionFieldType(fieldType) || this.isInterfaceFieldType(fieldType) || containsGenerationDirective(queryObj)) {
         return acc;
       }
       return { ...acc, [queryObj.name]: this.generateGraphQLOperationMetadata<CodeGenQuery, SchemaQuery>(queryObj) };
     }, {})
+    const generations = Object.values(this.queryMap).reduce((acc, queryObj: CodeGenQuery) => {
+      // Skip the field if the field type is union/interface
+      // TODO: Remove this skip once these types are supported for stakeholder usages
+      const fieldType = this.getType(queryObj.type) as any;
+      if (this.isUnionFieldType(fieldType) || this.isInterfaceFieldType(fieldType) || !containsGenerationDirective(queryObj)) {
+        return acc;
+      }
+      return { ...acc, [queryObj.name]: this.generateGenerationMetadata(queryObj) };
+    }, {});
     const mutations = Object.values(this.mutationMap).reduce((acc, mutationObj: CodeGenMutation) => {
       // Skip the field if the field type is union/interface
       // TODO: Remove this skip once these types are supported for stakeholder usages
@@ -95,6 +106,9 @@ export class AppSyncModelIntrospectionVisitor<
     if (Object.keys(queries).length > 0) {
       result = { ...result, queries };
     }
+    if (Object.keys(generations).length > 0) {
+      result = { ...result, generations };
+    }
     if (Object.keys(mutations).length > 0) {
       result = { ...result, mutations };
     }
@@ -236,6 +250,10 @@ export class AppSyncModelIntrospectionVisitor<
     return operationMeta as V;
   }
 
+  private generateGenerationMetadata(generationObj: CodeGenQuery): SchemaQuery {
+    return this.generateGraphQLOperationMetadata<CodeGenQuery, SchemaQuery>(generationObj);
+  }
+
   protected getType(gqlType: string): FieldType | InputFieldType | UnionFieldType | InterfaceFieldType {
     // Todo: Handle unlisted scalars
     if (gqlType in METADATA_SCALAR_MAP) {

Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,9 @@`
`26`	`26`	`},`
`27`	`27`	`"inputs": {`
`28`	`28`	`"$ref": "#/definitions/SchemaInputs"`
	`29`	`+ },`
	`30`	`+ "generations": {`
	`31`	`+ "$ref": "#/definitions/SchemaGenerations"`
`29`	`32`	`}`
`30`	`33`	`},`
`31`	`34`	`"required": [`
`@@ -532,6 +535,9 @@`
`532`	`535`	`],`
`533`	`536`	`"additionalProperties": false,`
`534`	`537`	`"description": "Input Definition"`
	`538`	`+ },`
	`539`	`+ "SchemaGenerations": {`
	`540`	`+ "$ref": "#/definitions/SchemaQueries"`
`535`	`541`	`}`
`536`	`542`	`}`
`537`	`543`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { CodeGenDirective, CodeGenField, CodeGenModel } from '../visitors/appsync-visitor';`
	`1`	`+import { CodeGenDirective, CodeGenField, CodeGenModel, CodeGenQuery } from '../visitors/appsync-visitor';`
`2`	`2`	`import { TransformerV2DirectiveName } from './constants';`
`3`	`3`
`4`	`4`	`export function addFieldToModel(model: CodeGenModel, field: CodeGenField): void {`
`@@ -40,4 +40,8 @@ export function getModelPrimaryKeyComponentFields(model: CodeGenModel): CodeGenF`
`40`	`40`	`};`
`41`	`41`	`}`
`42`	`42`	`return keyFields;`
	`43`	`+}`
	`44`	`+`
	`45`	`+export function containsGenerationDirective(queryField: CodeGenQuery): boolean {`
	`46`	`+ return queryField.directives.some((directive) => directive.name === TransformerV2DirectiveName.GENERATION);`
`43`	`47`	`}`