@aws-amplify/data-schema 1.17.5 → 1.19.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws-amplify/data-schema",
-  "version": "1.17.5",
+  "version": "1.19.0",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",

package/src/ClientSchema/Core/ClientCustomOperations.ts

@@ -9,8 +9,8 @@ import type { AppSyncResolverHandler } from 'aws-lambda';
 import type { CustomType } from '../../CustomType';
 import type { FieldTypesOfCustomType } from '../../MappedTypes/ResolveSchema';
 import type { ResolveRef } from '../utilities/ResolveRef';
-import type { EnumType } from '../../EnumType';
 import { ClientSchemaProperty } from './ClientSchemaProperty';
+import type { ResolveFields } from '../utilities';
 
 type CustomOperationSubType<Op extends CustomOperationParamShape> =
   `custom${Op['typeName']}`;
@@ -38,7 +38,7 @@ export interface ClientCustomOperation<
    * ```
    */
   functionHandler: AppSyncResolverHandler<
-    CustomOpArguments<Op>,
+    CustomOpArguments<Op, RefBag>,
     // If the final handler is an async function, the Schema['fieldname']['functionhandler']
     // should have a return type of `void`. This only applies to `functionHandler` and not
     // `returnType` because `returnType` determines the type returned by the mutation / query
@@ -60,7 +60,7 @@ export interface ClientCustomOperation<
    * }
    * ```
    */
-  args: CustomOpArguments<Op>;
+  args: CustomOpArguments<Op, RefBag>;
 
   /**
    * The return type expected by a lambda function handler.
@@ -84,19 +84,18 @@ export interface ClientCustomOperation<
 
 /**
  * Digs out custom operation arguments, mapped to the intended graphql types.
+ * using the existing ResolveFields utility type. This handles:
+ * - Basic scalar fields
+ * - Enum types
+ * - Custom types (including nested structures)
+ * - Reference types
  */
-type CustomOpArguments<Shape extends CustomOperationParamShape> =
-  Shape['arguments'] extends null
-    ? never
-    : ResolveFieldRequirements<{
-        [FieldName in keyof Shape['arguments']]: Shape['arguments'][FieldName] extends BaseModelField<
-          infer R
-        >
-          ? R
-          : Shape['arguments'][FieldName] extends EnumType<infer Values>
-            ? Values[number] | null
-            : never;
-      }>;
+type CustomOpArguments<
+  Shape extends CustomOperationParamShape,
+  RefBag extends Record<string, any> = any,
+> = Shape['arguments'] extends null
+  ? never
+  : ResolveFields<RefBag, Shape['arguments']>;
 
 /**
  * Removes `null | undefined` from the return type if the operation is a subscription,

package/src/CustomOperation.ts

@@ -29,7 +29,7 @@ type CustomOperationBrand =
   | typeof subscriptionBrand
   | typeof generationBrand;
 
-type CustomArguments = Record<string, BaseModelField | EnumType>;
+type CustomArguments = Record<string, BaseModelField | EnumType | CustomType<any> | RefType<any, any>>;
 type SubscriptionSource = RefType<any, any>;
 type InternalSubscriptionSource = InternalRef;
 type CustomReturnType = RefType<any> | CustomType<any>;

package/src/SchemaProcessor.ts

@@ -347,6 +347,7 @@ function customOperationToGql(
   customTypeAuthRules: CustomTypeAuthRules;
   lambdaFunctionDefinition: LambdaFunctionDefinition;
   customSqlDataSourceStrategy: CustomSqlDataSourceStrategy | undefined;
+  inputTypes: { name: string; fields: Record<string, any> }[];
 } {
   const {
     arguments: fieldArgs,
@@ -442,15 +443,21 @@ function customOperationToGql(
     });
   }
 
-  if (Object.keys(fieldArgs).length > 0) {
-    const { gqlFields, implicitTypes: implied } = processFields(
-      typeName,
-      fieldArgs,
-      {},
-      {},
-    );
-    callSignature += `(${gqlFields.join(', ')})`;
-    implicitTypes.push(...implied);
+  const { inputTypes, argDefinitions, collectedEnums } = generateInputTypes(
+    typeName,
+    fieldArgs,
+    getRefType,
+  );
+
+  // Handle collected enums
+  for (const [enumName, enumDef] of collectedEnums) {
+    if (!implicitTypes.some(([name]) => name === enumName)) {
+      implicitTypes.push([enumName, enumDef]);
+    }
+  }
+
+  if (argDefinitions.length > 0) {
+    callSignature += `(${argDefinitions.join(', ')})`;
   }
 
   const handler = handlers && handlers[0];
@@ -550,6 +557,7 @@ function customOperationToGql(
     customTypeAuthRules,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
+    inputTypes,
   };
 }
 
@@ -1316,6 +1324,174 @@ const mergeCustomTypeAuthRules = (
   }
 };
 
+/**
+ * Generates input types for custom operations in the schema.
+ *
+ * Processes operation arguments to create corresponding input types,
+ * handling referenced and inline custom types, enums, and nested structures.
+ * Manages circular references and prevents duplicate processing.
+ *
+ **/
+
+function generateInputTypes(
+  operationName: string,
+  args: Record<string, any>,
+  getRefType: (name: string, referrerName?: string) => GetRef,
+): {
+  inputTypes: { name: string; fields: Record<string, any> }[];
+  argDefinitions: string[];
+  collectedEnums: Map<string, any>;
+} {
+  const inputTypes: {
+    name: string;
+    fields: Record<string, any>;
+  }[] = [];
+  const argDefinitions: string[] = [];
+  const collectedEnums: Map<string, any> = new Map();
+  const processedTypes = new Set<string>(); // Track processed types to avoid duplicates
+
+  const processNonScalarFields = (
+    fields: Record<string, any>,
+    originalTypeName: string,
+    isParentRef: boolean = false,
+    parentChain: string[] = [], // Used to detect circular references
+  ): Record<string, any> => {
+    const processedFields: Record<string, any> = {};
+    for (const [fieldName, fieldDef] of Object.entries(fields)) {
+      if (isRefField(fieldDef)) {
+        const refType = getRefType(fieldDef.data.link, originalTypeName);
+        if (refType.type === 'CustomType') {
+          const nestedInputTypeName = `${fieldDef.data.link}Input`;
+          processedFields[fieldName] = {
+            data: { type: 'ref', link: nestedInputTypeName },
+          };
+
+          // Process the nested type if it hasn't been processed and isn't a circular reference
+          if (
+            !parentChain.includes(nestedInputTypeName) &&
+            !processedTypes.has(nestedInputTypeName)
+          ) {
+            processedTypes.add(nestedInputTypeName);
+            const nestedFields = processNonScalarFields(
+              refType.def.data.fields,
+              fieldDef.data.link,
+              true,
+              [...parentChain, nestedInputTypeName],
+            );
+            inputTypes.push({
+              name: nestedInputTypeName,
+              fields: nestedFields,
+            });
+          }
+        } else if (refType.type === 'Enum') {
+          processedFields[fieldName] = {
+            data: { type: 'ref', link: fieldDef.data.link },
+          };
+        } else {
+          throw new Error(
+            `Unsupported reference type '${refType.type}' for field '${fieldName}'. ` +
+              `Only references to CustomType and Enum are supported.`,
+          );
+        }
+      } else if (isCustomType(fieldDef)) {
+        // Handle inline custom types
+        const nestedInputTypeName = `${capitalize(originalTypeName)}${capitalize(fieldName)}Input`;
+        processedFields[fieldName] = {
+          data: { type: 'ref', link: nestedInputTypeName },
+        };
+
+        if (!processedTypes.has(nestedInputTypeName)) {
+          processedTypes.add(nestedInputTypeName);
+          const nestedFields = processNonScalarFields(
+            fieldDef.data.fields,
+            `${capitalize(originalTypeName)}${capitalize(fieldName)}`,
+            isParentRef,
+            [...parentChain, nestedInputTypeName],
+          );
+          inputTypes.push({
+            name: nestedInputTypeName,
+            fields: nestedFields,
+          });
+        }
+      } else if (isEnumType(fieldDef)) {
+        // Handle enum types
+        const enumName = `${capitalize(originalTypeName)}${capitalize(fieldName)}`;
+        if (!collectedEnums.has(enumName) && !isParentRef) {
+          collectedEnums.set(enumName, fieldDef);
+        }
+        processedFields[fieldName] = { data: { type: 'ref', link: enumName } };
+      } else {
+        processedFields[fieldName] = fieldDef;
+      }
+    }
+    return processedFields;
+  };
+
+  // Process top-level arguments
+  for (const [argName, argDef] of Object.entries(args)) {
+    if (isRefField(argDef)) {
+      const refType = getRefType(argDef.data.link, operationName);
+      if (refType.type === 'CustomType') {
+        const inputTypeName = `${argDef.data.link}Input`;
+        argDefinitions.push(`${argName}: ${inputTypeName}`);
+
+        // Process the input type if it hasn't been processed yet
+        if (!processedTypes.has(inputTypeName)) {
+          processedTypes.add(inputTypeName);
+          const fields = processNonScalarFields(
+            refType.def.data.fields,
+            argDef.data.link,
+            true,
+            [inputTypeName],
+          );
+          inputTypes.push({
+            name: inputTypeName,
+            fields,
+          });
+        }
+      } else if (refType.type === 'Enum') {
+        argDefinitions.push(`${argName}: ${argDef.data.link}`);
+      } else {
+        throw new Error(
+          `Unsupported reference type '${refType.type}' for argument '${argName}' in '${operationName}'. ` +
+            `Only references to CustomType and Enum are supported.`,
+        );
+      }
+    } else if (isEnumType(argDef)) {
+      // Handle top-level enum arguments
+      const enumName = `${capitalize(operationName)}${capitalize(argName)}`;
+      if (!collectedEnums.has(enumName)) {
+        collectedEnums.set(enumName, argDef);
+      }
+      argDefinitions.push(`${argName}: ${enumName}`);
+    } else if (isCustomType(argDef)) {
+      // Handle top-level custom type arguments
+      const inputTypeName = `${capitalize(operationName)}${capitalize(argName)}Input`;
+      argDefinitions.push(`${argName}: ${inputTypeName}`);
+
+      if (!processedTypes.has(inputTypeName)) {
+        processedTypes.add(inputTypeName);
+        const fields = processNonScalarFields(
+          argDef.data.fields,
+          `${capitalize(operationName)}${capitalize(argName)}`,
+          false,
+          [inputTypeName],
+        );
+        inputTypes.push({
+          name: inputTypeName,
+          fields,
+        });
+      }
+    } else if (isScalarField(argDef)) {
+      argDefinitions.push(`${argName}: ${scalarFieldToGql(argDef.data)}`);
+    } else {
+      throw new Error(`Unsupported argument type for ${argName}`);
+    }
+  }
+
+  return { inputTypes, argDefinitions, collectedEnums };
+}
+
 const schemaPreprocessor = (
   schema: InternalSchema,
 ): {
@@ -1326,7 +1502,6 @@ const schemaPreprocessor = (
   customSqlDataSourceStrategies?: CustomSqlDataSourceStrategy[];
 } => {
   const gqlModels: string[] = [];
-
   const customQueries = [];
   const customMutations = [];
   const customSubscriptions = [];
@@ -1374,6 +1549,7 @@ const schemaPreprocessor = (
   );
 
   const getRefType = getRefTypeForSchema(schema);
+  const uniqueInputTypes = new Map<string, Record<string, any>>();
 
   for (const [typeName, typeDef] of topLevelTypes) {
     const mostRelevantAuthRules: Authorization<any, any, any>[] =
@@ -1449,6 +1625,7 @@ const schemaPreprocessor = (
           jsFunctionForField,
           lambdaFunctionDefinition,
           customSqlDataSourceStrategy,
+          inputTypes,
         } = transformCustomOperations(
           typeDef,
           typeName,
@@ -1457,6 +1634,13 @@ const schemaPreprocessor = (
           getRefType,
         );
 
+        // Process input types without duplicates
+        for (const { name, fields } of inputTypes) {
+          if (!uniqueInputTypes.has(name)) {
+            uniqueInputTypes.set(name, fields);
+          }
+        }
+
         topLevelTypes.push(...implicitTypes);
 
         mergeCustomTypeAuthRules(
@@ -1542,7 +1726,6 @@ const schemaPreprocessor = (
         undefined,
         databaseEngine,
       );
-
       topLevelTypes.push(...implicitTypes);
 
       const joined = gqlFields.join('\n  ');
@@ -1622,6 +1805,22 @@ const schemaPreprocessor = (
     }
   }
 
+  // Generate input types after processing all custom operations
+  for (const [name, fields] of uniqueInputTypes) {
+    const { gqlFields } = processFields(
+      name,
+      fields,
+      {},
+      {},
+      undefined,
+      undefined,
+      undefined,
+      databaseEngine,
+    );
+    const inputTypeDefinition = `input ${name} {\n  ${gqlFields.join('\n  ')}\n}`;
+    gqlModels.push(inputTypeDefinition);
+  }
+
   const customOperations = {
     queries: customQueries,
     mutations: customMutations,
@@ -1632,7 +1831,6 @@ const schemaPreprocessor = (
   if (shouldAddConversationTypes) {
     gqlModels.push(CONVERSATION_SCHEMA_GRAPHQL_TYPES);
   }
-
   const processedSchema = gqlModels.join('\n\n');
 
   return {
@@ -1935,6 +2133,7 @@ function transformCustomOperations(
     customTypeAuthRules,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
+    inputTypes,
   } = customOperationToGql(
     typeName,
     typeDef,
@@ -1951,6 +2150,7 @@ function transformCustomOperations(
     jsFunctionForField,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
+    inputTypes,
   };
 }
 

package/src/ai/types/ConversationMessageContent.ts

@@ -1,11 +1,17 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { ImageBlock, ToolResultBlock, ToolUseBlock } from './contentBlocks';
+import {
+  ImageBlock,
+  DocumentBlock,
+  ToolResultBlock,
+  ToolUseBlock,
+} from './contentBlocks';
 
 export interface ConversationMessageTextContent {
   text: string;
   image?: never;
+  document?: never;
   toolUse?: never;
   toolResult?: never;
 }
@@ -13,13 +19,23 @@ export interface ConversationMessageTextContent {
 export interface ConversationMessageImageContent {
   text?: never;
   image: ImageBlock;
+  document?: never;
   toolUse?: never;
   toolResult?: never;
 }
 
+export interface ConversationMessageDocumentContent {
+  text?: never;
+  image?: never;
+  toolUse?: never;
+  toolResult?: never;
+  document: DocumentBlock;
+}
+
 export interface ConversationMessageToolUseContent {
   text?: never;
   image?: never;
+  document?: never;
   toolUse: ToolUseBlock;
   toolResult?: never;
 }
@@ -27,6 +43,7 @@ export interface ConversationMessageToolUseContent {
 export interface ConversationMessageToolResultContent {
   text?: never;
   image?: never;
+  document?: never;
   toolUse?: never;
   toolResult: ToolResultBlock;
 }
@@ -34,10 +51,12 @@ export interface ConversationMessageToolResultContent {
 export type ConversationMessageContent =
   | ConversationMessageTextContent
   | ConversationMessageImageContent
+  | ConversationMessageDocumentContent
   | ConversationMessageToolUseContent
   | ConversationMessageToolResultContent;
 
 export type ConversationSendMessageInputContent =
   | Omit<ConversationMessageTextContent, 'toolUse'>
   | Omit<ConversationMessageImageContent, 'toolUse'>
+  | Omit<ConversationMessageDocumentContent, 'toolUse'>
   | Omit<ConversationMessageToolResultContent, 'toolUse'>;

package/src/ai/types/contentBlocks.ts

@@ -8,12 +8,31 @@ interface BytesImageSource {
   bytes: Uint8Array;
 }
 
+interface BytesDocumentSource {
+  bytes: Uint8Array;
+}
+
 // common content blocks
 export interface ImageBlock {
   format: 'gif' | 'jpeg' | 'png' | 'webp';
   source: BytesImageSource;
 }
 
+export interface DocumentBlock {
+  format:
+    | 'pdf'
+    | 'csv'
+    | 'doc'
+    | 'docx'
+    | 'xls'
+    | 'xlsx'
+    | 'html'
+    | 'txt'
+    | 'md';
+  name: string;
+  source: BytesDocumentSource;
+}
+
 export interface ToolUseBlock {
   toolUseId: string;
   name: string;

package/src/runtime/bridge-types.ts

@@ -16,6 +16,8 @@ import { Observable } from 'rxjs';
 import { CustomHeaders, ModelSortDirection } from './client';
 
 import { AiAction, AiCategory } from './internals/ai/getCustomUserAgentDetails';
+import { CustomType } from '../CustomType';
+import { RefType } from '../RefType';
 
 export declare namespace AmplifyServer {
   export interface ContextToken {
@@ -240,7 +242,12 @@ export type FieldType =
   | ModelFieldType
   | NonModelFieldType;
 
-export type InputFieldType = ScalarType | EnumType | InputType;
+export type InputFieldType =
+  | ScalarType
+  | EnumType
+  | InputType
+  | CustomType<any>
+  | RefType<any, any>;
 
 export type FieldAttribute = ModelAttribute;
 

package/src/runtime/internals/ai/conversationMessageDeserializers.ts

@@ -11,6 +11,9 @@ export const deserializeContent = (
     if (block.image) {
       return deserializeImageBlock(block);
     }
+    if (block.document) {
+      return deserializeDocumentBlock(block);
+    }
     if (block.toolUse) {
       return deserializeToolUseBlock(block);
     }
@@ -30,6 +33,16 @@ const deserializeImageBlock = ({ image }: Record<'image', any>) => ({
   },
 });
 
+const deserializeDocumentBlock = ({ document }: Record<'document', any>) => ({
+  document: {
+    ...document,
+    source: {
+      ...document.source,
+      bytes: fromBase64(document.source.bytes),
+    },
+  },
+});
+
 const deserializeJsonBlock = ({ json }: Record<'json', any>) => ({
   json: JSON.parse(json),
 });

package/src/runtime/internals/ai/conversationMessageSerializers.ts

@@ -5,6 +5,7 @@ import { toBase64 } from '@smithy/util-base64';
 import type {
   ConversationMessageContent,
   ConversationMessageImageContent,
+  ConversationMessageDocumentContent,
   ConversationMessageToolResultContent,
 } from '../../../ai/types/ConversationMessageContent';
 import type { ToolConfiguration } from '../../../ai/types/ToolConfiguration';
@@ -18,6 +19,9 @@ export const serializeContent = (content: ConversationMessageContent[]) =>
     if (block.image) {
       return serializeImageBlock(block);
     }
+    if (block.document) {
+      return serializeDocumentBlock(block);
+    }
     if (block.toolResult) {
       return serializeToolResultBlock(block);
     }
@@ -46,6 +50,18 @@ const serializeImageBlock = ({ image }: ConversationMessageImageContent) => ({
   },
 });
 
+const serializeDocumentBlock = ({
+  document,
+}: ConversationMessageDocumentContent) => ({
+  document: {
+    ...document,
+    source: {
+      ...document.source,
+      bytes: toBase64(document.source.bytes),
+    },
+  },
+});
+
 const serializeJsonBlock = ({ json }: ToolResultJsonContent) => ({
   json: JSON.stringify(json),
 });