@aws-amplify/data-schema 1.4.0 → 1.5.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws-amplify/data-schema",
-  "version": "1.4.0",
+  "version": "1.5.0",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",
@@ -51,7 +51,9 @@
   },
   "dependencies": {
     "@aws-amplify/data-schema-types": "*",
+    "@smithy/util-base64": "^3.0.0",
     "@types/aws-lambda": "^8.10.134",
+    "@types/json-schema": "^7.0.15",
     "rxjs": "^7.8.1"
   },
   "devDependencies": {

package/src/ClientSchema/Core/ClientSchemaProperty.ts

@@ -5,6 +5,8 @@ export interface ClientSchemaProperty {
     | 'customType'
     | 'customQuery'
     | 'customMutation'
-    | 'customSubscription';
+    | 'customSubscription'
+    | 'customConversation'
+    | 'customGeneration';
   type: any;
 }

package/src/ClientSchema/ai/ClientConversation.ts

@@ -0,0 +1,15 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type {
+  Conversation,
+  ConversationMessage,
+} from '../../ai/ConversationType';
+import type { ClientSchemaProperty } from '../Core';
+
+export interface ClientConversation
+  extends Pick<ClientSchemaProperty, '__entityType'> {
+  __entityType: 'customConversation';
+  type: Conversation;
+  messageType: ConversationMessage;
+}

package/src/ClientSchema/index.ts

@@ -1,6 +1,7 @@
 import type { __modelMeta__ } from '../runtime/client';
 import type {
   BaseSchema,
+  ConversationType,
   CustomOperation,
   CustomType,
   EnumType,
@@ -21,6 +22,7 @@ import type {
 } from '../CombineSchema';
 import type { SchemaMetadata } from './utilities/SchemaMetadata';
 import type { Brand, Select, SpreadTuple } from '../util';
+import { ClientConversation } from './ai/ClientConversation';
 
 export type ClientSchema<
   Schema extends GenericModelSchema<any> | CombinedModelSchema<any>,
@@ -65,11 +67,14 @@ type ClientSchemaProperty<
             | 'queryCustomOperation'
             | 'mutationCustomOperation'
             | 'subscriptionCustomOperation'
+            | 'generationCustomOperation'
           >
         ? RemapCustomOperation<T, Metadata, IsRDS, T[K]>
         : T[K] extends Brand<'modelType'>
           ? RemapModel<T, Metadata, IsRDS, T[K], K>
-          : never;
+          : T[K] extends Brand<'conversationCustomOperation'>
+            ? RemapAIRoute<T, T[K]>
+            : never;
 
 type RemapEnum<_T extends ModelSchemaContents, E> =
   E extends EnumType<infer values> ? ClientEnum<values> : never;
@@ -111,6 +116,11 @@ type RemapModel<
       >
     : never;
 
+type RemapAIRoute<
+  _T extends ModelSchemaContents,
+  E,
+> = E extends ConversationType ? ClientConversation : never;
+
 type GetInternalClientSchema<Schema> =
   Schema extends GenericModelSchema<any> ? InternalClientSchema<Schema> : never;
 
@@ -146,6 +156,8 @@ export type ClientSchemaByEntityTypeBaseShape = {
   queries: Record<string, ClientCustomOperation<any, any>>;
   mutations: Record<string, ClientCustomOperation<any, any>>;
   subscriptions: Record<string, ClientCustomOperation<any, any>>;
+  conversations: Record<string, ClientConversation>;
+  generations: Record<string, ClientCustomOperation<any, any>>;
 };
 
 export type ClientSchemaByEntityType<T> = {
@@ -155,4 +167,6 @@ export type ClientSchemaByEntityType<T> = {
   queries: Select<T, { __entityType: 'customQuery' }>;
   mutations: Select<T, { __entityType: 'customMutation' }>;
   subscriptions: Select<T, { __entityType: 'customSubscription' }>;
+  conversations: Select<T, { __entityType: 'customConversation' }>;
+  generations: Select<T, { __entityType: 'customGeneration' }>;
 };

package/src/CustomOperation.ts

@@ -14,15 +14,18 @@ import type {
   FunctionHandler,
   HandlerType as Handler,
 } from './Handler';
+import { AiModel, InferenceConfiguration } from './ai/ModelType';
 
 const queryBrand = 'queryCustomOperation';
 const mutationBrand = 'mutationCustomOperation';
 const subscriptionBrand = 'subscriptionCustomOperation';
+const generationBrand = 'generationCustomOperation';
 
 type CustomOperationBrand =
   | typeof queryBrand
   | typeof mutationBrand
-  | typeof subscriptionBrand;
+  | typeof subscriptionBrand
+  | typeof generationBrand;
 
 type CustomArguments = Record<string, BaseModelField | EnumType>;
 
@@ -38,6 +41,7 @@ export const CustomOperationNames = [
   'Query',
   'Mutation',
   'Subscription',
+  'Generation',
 ] as const;
 type CustomOperationName = (typeof CustomOperationNames)[number];
 
@@ -48,6 +52,7 @@ type CustomData = {
   typeName: CustomOperationName;
   handlers: Handler[] | null;
   subscriptionSource: SubscriptionSource[];
+  input?: CustomOperationInput;
 };
 
 type InternalCustomData = CustomData & {
@@ -57,12 +62,15 @@ type InternalCustomData = CustomData & {
   authorization: Authorization<any, any, any>[];
 };
 
+export type CustomOperationInput = GenerationInput;
+
 export type CustomOperationParamShape = {
   arguments: CustomArguments | null;
   returnType: CustomReturnType | null;
   authorization: Authorization<any, any, any>[];
   typeName: CustomOperationName;
   handlers: Handler | null;
+  input?: CustomOperationInput;
 };
 
 export type CustomOperation<
@@ -137,7 +145,7 @@ export type InternalCustom<B extends CustomOperationBrand = any> =
 function _custom<
   T extends CustomOperationParamShape,
   B extends CustomOperationBrand,
->(typeName: CustomOperationName, brand: B) {
+>(typeName: CustomOperationName, brand: B, input?: T['input']) {
   const data: CustomData = {
     arguments: {},
     returnType: null,
@@ -145,6 +153,7 @@ function _custom<
     typeName: typeName,
     handlers: null,
     subscriptionSource: [],
+    input,
   };
 
   const builder = brandedBuilder<T>(
@@ -210,7 +219,7 @@ export type QueryCustomOperation = CustomOperation<
  *     .authorization(allow => [allow.publicApiKey()])
  *     // 3. set the function has the handler
  *     .handler(a.handler.function(echoHandler)),
- * 
+ *
  *   EchoResponse: a.customType({
  *     content: a.string(),
  *     executionDuration: a.float()
@@ -277,9 +286,9 @@ export type SubscriptionCustomOperation = CustomOperation<
  * // Subscribe to incoming messages
  * receive: a.subscription()
  *   // subscribes to the 'publish' mutation
- *   .for(a.ref('publish')) 
+ *   .for(a.ref('publish'))
  *   // subscription handler to set custom filters
- *   .handler(a.handler.custom({entry: './receive.js'})) 
+ *   .handler(a.handler.custom({entry: './receive.js'}))
  *   // authorization rules as to who can subscribe to the data
  *   .authorization(allow => [allow.publicApiKey()]),
  * @returns a custom subscription
@@ -297,3 +306,37 @@ export function subscription(): CustomOperation<
 > {
   return _custom('Subscription', subscriptionBrand);
 }
+
+export interface GenerationInput {
+  aiModel: AiModel;
+  systemPrompt: string;
+  inferenceConfiguration?: InferenceConfiguration;
+}
+
+/**
+ * @experimental
+ *
+ * Define an AI generation route for single request-response interaction with specified AI model.
+ * @example
+ * makeRecipe: a.generation({
+ *   aiModel: { resourcePath },
+ *   systemPrompt: 'Please make a recipe from the provided ingredients',
+ * })
+ *   .arguments({ ingredients: a.string().array() })
+ *   .returns(a.ref("Recipe"))
+ * @returns a generation route definition
+ */
+export function generation(input: GenerationInput): CustomOperation<
+  {
+    arguments: null;
+    returnType: null;
+    authorization: [];
+    typeName: 'Generation';
+    handlers: null;
+    input: GenerationInput;
+  },
+  'for' | 'handler',
+  typeof generationBrand
+> {
+  return _custom('Generation', generationBrand, input);
+}

package/src/ModelSchema.ts

@@ -30,11 +30,13 @@ import {
   ModelRelationalField,
   ModelRelationalFieldParamShape,
 } from './ModelRelationalField';
+import { ConversationType } from './ai/ConversationType';
 
 export { ModelType } from './ModelType';
 export { EnumType } from './EnumType';
 export { CustomType } from './CustomType';
 export { CustomOperation } from './CustomOperation';
+export { ConversationType } from './ai/ConversationType';
 
 export const rdsSchemaBrandName = 'RDSSchema';
 export const rdsSchemaBrand = brand(rdsSchemaBrandName);
@@ -48,7 +50,8 @@ type SchemaContent =
   | BaseModelType
   | CustomType<CustomTypeParamShape>
   | EnumType
-  | CustomOperation<CustomOperationParamShape, any>;
+  | CustomOperation<CustomOperationParamShape, any>
+  | ConversationType;
 
 // The SQL-only `addToSchema` accepts all top-level entities, excepts models
 type AddToSchemaContent = Exclude<SchemaContent, BaseModelType>;

package/src/SchemaProcessor.ts

@@ -1,4 +1,4 @@
-import { type CustomPathData, type InternalSchema } from './ModelSchema';
+import type { CustomPathData, InternalSchema } from './ModelSchema';
 import {
   type ModelField,
   type InternalField,
@@ -30,17 +30,30 @@ import {
 import type { InternalRef, RefType } from './RefType';
 import type { EnumType } from './EnumType';
 import type { CustomType, CustomTypeParamShape } from './CustomType';
-import { type InternalCustom, CustomOperationNames } from './CustomOperation';
+import {
+  type InternalCustom,
+  type CustomOperationInput,
+  type GenerationInput,
+  CustomOperationNames,
+} from './CustomOperation';
 import { Brand, getBrand } from './util';
 import {
   getHandlerData,
   type HandlerType,
   type CustomHandler,
   type SqlReferenceHandler,
-  FunctionHandler,
+  type FunctionHandler,
 } from './Handler';
 import * as os from 'os';
 import * as path from 'path';
+import {
+  brandName as conversationBrandName,
+  type InternalConversationType,
+} from './ai/ConversationType';
+import {
+  conversationTypes,
+  createConversationField,
+} from './ai/ConversationSchemaTypes';
 
 type ScalarFieldDef = Exclude<InternalField['data'], { fieldType: 'model' }>;
 
@@ -84,6 +97,16 @@ function isCustomType(
   return false;
 }
 
+function isConversationRoute(type: any): type is InternalConversationType {
+  return getBrand(type) === conversationBrandName;
+}
+
+function isGenerationInput(
+  input?: CustomOperationInput,
+): input is GenerationInput {
+  return Boolean(input?.aiModel && input?.systemPrompt);
+}
+
 function isCustomOperation(type: any): type is InternalCustom {
   if (CustomOperationNames.includes(type?.data?.typeName)) {
     return true;
@@ -469,7 +492,29 @@ function customOperationToGql(
     gqlHandlerContent += `@aws_subscribe(mutations: ["${subscriptionSources}"]) `;
   }
 
+  if (opType === 'Generation') {
+    if (!isGenerationInput(typeDef.data.input)) {
+      throw new Error(
+        `Invalid Generation Route definition. A Generation Route must include a valid input. ${typeName} has an invalid or no input defined.`,
+      );
+    }
+    const { aiModel, systemPrompt, inferenceConfiguration } =
+      typeDef.data.input;
+
+    const inferenceConfigurationEntries = Object.entries(
+      inferenceConfiguration ?? {},
+    );
+    const inferenceConfigurationGql =
+      inferenceConfigurationEntries.length > 0
+        ? `, inferenceConfiguration: { ${inferenceConfigurationEntries
+            .map(([key, value]) => `${key}: ${value}`)
+            .join(', ')} }`
+        : '';
+    gqlHandlerContent += `@generation(aiModel: "${aiModel.resourcePath}", systemPrompt: "${systemPrompt}"${inferenceConfigurationGql}) `;
+  }
+
   const gqlField = `${callSignature}: ${returnTypeName} ${gqlHandlerContent}${authString}`;
+
   return {
     gqlField,
     implicitTypes: implicitTypes,
@@ -1210,6 +1255,7 @@ const schemaPreprocessor = (
   const customQueries = [];
   const customMutations = [];
   const customSubscriptions = [];
+  let shouldAddConversationTypes = false;
 
   // Dict of auth rules to be applied to custom types
   // Inherited from the auth configured on the custom operations that return these custom types
@@ -1296,6 +1342,8 @@ const schemaPreprocessor = (
         const model = `type ${typeName} ${customAuth}\n{\n  ${joined}\n}`;
         gqlModels.push(model);
       } else if (isCustomOperation(typeDef)) {
+        // TODO: add generation route logic.
+
         const { typeName: opType } = typeDef.data;
 
         const {
@@ -1347,6 +1395,7 @@ const schemaPreprocessor = (
 
         switch (opType) {
           case 'Query':
+          case 'Generation':
             customQueries.push(gqlField);
             break;
           case 'Mutation':
@@ -1358,6 +1407,10 @@ const schemaPreprocessor = (
           default:
             break;
         }
+      } else if (isConversationRoute(typeDef)) {
+        // TODO: add inferenceConfiguration values to directive.
+        customMutations.push(createConversationField(typeDef, typeName));
+        shouldAddConversationTypes = true;
       }
     } else if (staticSchema) {
       const fields = { ...typeDef.data.fields } as Record<
@@ -1473,6 +1526,9 @@ const schemaPreprocessor = (
   };
 
   gqlModels.push(...generateCustomOperationTypes(customOperations));
+  if (shouldAddConversationTypes) {
+    gqlModels.push(...conversationTypes);
+  }
 
   const processedSchema = gqlModels.join('\n\n');
 
@@ -1497,8 +1553,9 @@ function validateCustomOperations(
   const authConfigured = authRules.length > 0;
 
   if (
-    (authConfigured && !handlerConfigured) ||
-    (handlerConfigured && !authConfigured)
+    opType !== 'Generation' &&
+    ((authConfigured && !handlerConfigured) ||
+      (handlerConfigured && !authConfigured))
   ) {
     // Deploying a custom operation with auth and no handler reference OR
     // with a handler reference but no auth
@@ -1529,10 +1586,12 @@ function validateCustomOperations(
 
   if (
     typeDef.data.returnType === null &&
-    (opType === 'Query' || opType === 'Mutation')
+    (opType === 'Query' || opType === 'Mutation' || opType === 'Generation')
   ) {
+    const typeDescription =
+      opType === 'Generation' ? 'Generation Route' : `Custom ${opType}`;
     throw new Error(
-      `Invalid Custom ${opType} definition. A Custom ${opType} must include a return type. ${typeName} has no return type specified.`,
+      `Invalid ${typeDescription} definition. A ${typeDescription} must include a return type. ${typeName} has no return type specified.`,
     );
   }
 
@@ -1711,12 +1770,18 @@ function transformCustomOperations(
   getRefType: ReturnType<typeof getRefTypeForSchema>,
 ) {
   const { typeName: opType, handlers } = typeDef.data;
+
   let jsFunctionForField: JsResolver | undefined = undefined;
 
   validateCustomOperations(typeDef, typeName, authRules, getRefType);
 
   if (isCustomHandler(handlers)) {
-    jsFunctionForField = handleCustom(handlers, opType, typeName);
+    jsFunctionForField = handleCustom(
+      handlers,
+      // Generation routes should not have handlers
+      opType as Exclude<typeof opType, 'Generation'>,
+      typeName,
+    );
   }
 
   const isCustom = Boolean(jsFunctionForField);

package/src/a.ts

@@ -21,8 +21,14 @@ import { ref } from './RefType';
 import { hasOne, hasMany, belongsTo } from './ModelRelationalField';
 import { customType } from './CustomType';
 import { enumType } from './EnumType';
-import { query, mutation, subscription } from './CustomOperation';
+import { query, mutation, subscription, generation } from './CustomOperation';
 import { handler } from './Handler';
+import { conversation } from './ai/ConversationType';
+import { model as aiModel } from './ai/ModelType';
+
+const ai = {
+  model: aiModel,
+};
 
 // We are re-exporting with this pattern to make the JSDoc happy.
 // Otherwise the JSDoc wouldn't consistently show up as a customer types
@@ -54,5 +60,8 @@ export {
   phone,
   url,
   ipAddress,
+  conversation,
+  generation,
+  ai,
   handler,
 };