@aws-amplify/data-schema 1.4.0 → 1.5.0

package/src/runtime/client/index.ts

@@ -8,6 +8,7 @@ import {
   __modelMeta__,
 } from '@aws-amplify/data-schema-types';
 import type { Observable } from 'rxjs';
+import type { ConversationRoute } from '../../ai/ConversationType';
 import type {
   ClientSchemaByEntityType,
   ClientSchemaByEntityTypeBaseShape,
@@ -511,31 +512,45 @@ type ModelTypesClient<
   // Omit any disabled operations
   Omit<
     {
-      create: (
+      create<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         model: Model['createType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-      update: (
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
+      update<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         model: Model['updateType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-      delete: (
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
+      delete<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         identifier: Model['deleteType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
       get<SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[]>(
         identifier: Model['identifier'],
         options?: {
@@ -558,7 +573,6 @@ type ModelTypesClient<
           headers?: CustomHeaders;
         },
       ): ListReturnValue<Prettify<ReturnValue<Model, FlatModel, SelectionSet>>>;
-
       onCreate<
         SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
       >(options?: {
@@ -610,33 +624,48 @@ type ModelTypesSSRCookies<
   Model extends ClientSchemaByEntityTypeBaseShape['models'][string],
   FlatModel extends Record<string, unknown> = ResolvedModel<Model['type']>,
 > = IndexQueryMethods<Model> &
+  // Omit any disabled operations
   Omit<
     {
-      create: (
+      create<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         model: Model['createType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-      update: (
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
+      update<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         model: Model['updateType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-      delete: (
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
+      delete<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         identifier: Model['deleteType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
       get<SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[]>(
         identifier: Model['identifier'],
         options?: {
@@ -668,36 +697,51 @@ type ModelTypesSSRRequest<
   Model extends ClientSchemaByEntityTypeBaseShape['models'][string],
   FlatModel extends Record<string, unknown> = ResolvedModel<Model['type']>,
 > = IndexQueryMethods<Model, 'REQUEST'> &
+  // Omit any disabled operations
   Omit<
     {
-      create: (
+      create<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         contextSpec: AmplifyServer.ContextSpec,
         model: Model['createType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-      update: (
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
+      update<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         contextSpec: AmplifyServer.ContextSpec,
         model: Model['updateType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-      delete: (
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
+      delete<
+        SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[],
+      >(
         contextSpec: AmplifyServer.ContextSpec,
         identifier: Model['deleteType'],
         options?: {
+          selectionSet?: SelectionSet;
           authMode?: AuthMode;
           authToken?: string;
           headers?: CustomHeaders;
         },
-      ) => SingularReturnValue<Model['type']>;
-
+      ): SingularReturnValue<
+        Prettify<ReturnValue<Model, FlatModel, SelectionSet>>
+      >;
       get<SelectionSet extends ReadonlyArray<ModelPath<FlatModel>> = never[]>(
         contextSpec: AmplifyServer.ContextSpec,
         identifier: Model['identifier'],
@@ -781,7 +825,8 @@ export type CustomOperations<
   OperationDefs extends ClientSchemaByEntityTypeBaseShape[
     | 'queries'
     | 'mutations'
-    | 'subscriptions'],
+    | 'subscriptions'
+    | 'generations'],
   Context extends ContextType = 'CLIENT',
 > = {
   [OpName in keyof OperationDefs]: {
@@ -857,6 +902,8 @@ export type ClientExtensions<T extends Record<any, any> = never> = {
   queries: CustomQueries<T, 'CLIENT'>;
   mutations: CustomMutations<T, 'CLIENT'>;
   subscriptions: CustomSubscriptions<T, 'CLIENT'>;
+  conversations: ConversationRoutes<T>;
+  generations: Generations<T>;
 };
 
 export type ClientExtensionsSSRRequest<T extends Record<any, any> = never> = {
@@ -872,3 +919,15 @@ export type ClientExtensionsSSRCookies<T extends Record<any, any> = never> = {
   queries: CustomQueries<T, 'COOKIES'>;
   mutations: CustomMutations<T, 'COOKIES'>;
 };
+
+export type ConversationRoutes<
+  T extends Record<any, any>,
+  Schema extends ClientSchemaByEntityType<T> = ClientSchemaByEntityType<T>,
+> = {
+  [ConversationName in keyof Schema['conversations']]: ConversationRoute;
+};
+
+export type Generations<
+  T extends Record<any, any>,
+  Schema extends ClientSchemaByEntityType<T> = ClientSchemaByEntityType<T>,
+> = CustomOperations<Schema['generations']>;

package/src/runtime/internals/APIClient.ts

@@ -248,7 +248,7 @@ export function initializeModel(
           );
 
           // if get is disabled on the related model
-          if ((client as any).models[relatedModelName].get === undefined) {
+          if ((client as any).models[relatedModelName]?.get === undefined) {
             break;
           }
 
@@ -340,7 +340,7 @@ export function initializeModel(
             );
 
             // if list is disabled on the related model
-            if ((client as any).models[relatedModelName].list === undefined) {
+            if ((client as any).models[relatedModelName]?.list === undefined) {
               break;
             }
 
@@ -411,7 +411,7 @@ export function initializeModel(
           );
 
           // if list is disabled on the related model
-          if ((client as any).models[relatedModelName].list === undefined) {
+          if ((client as any).models[relatedModelName]?.list === undefined) {
             break;
           }
 
@@ -1242,7 +1242,6 @@ export function buildGraphQLVariables(
       throw new Error(
         'Internal error: Attempted to build variables for observeQuery. Please report this error.',
       );
-      break;
     default: {
       const exhaustiveCheck: never = operation;
       throw new Error(`Unhandled operation case: ${exhaustiveCheck}`);

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

@@ -0,0 +1,62 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { fromBase64 } from '@smithy/util-base64';
+import type { ConversationMessageContent } from '../../../ai/types/ConversationMessageContent';
+
+export const deserializeContent = (
+  content: Record<string, any>[],
+): ConversationMessageContent[] =>
+  content.map((block) => {
+    if (block.image) {
+      return deserializeImageBlock(block);
+    }
+    if (block.toolUse) {
+      return deserializeToolUseBlock(block);
+    }
+    if (block.toolResult) {
+      return deserializeToolResultBlock(block);
+    }
+    return removeNullsFromBlock(block) as ConversationMessageContent;
+  });
+
+const deserializeImageBlock = ({ image }: Record<'image', any>) => ({
+  image: {
+    ...image,
+    source: {
+      ...image.source,
+      bytes: fromBase64(image.source.bytes),
+    },
+  },
+});
+
+const deserializeJsonBlock = ({ json }: Record<'json', any>) => ({
+  json: JSON.parse(json),
+});
+
+const deserializeToolUseBlock = ({ toolUse }: Record<'toolUse', any>) => ({
+  toolUse: {
+    ...toolUse,
+    input: JSON.parse(toolUse.input),
+  },
+});
+
+const deserializeToolResultBlock = ({
+  toolResult,
+}: Record<'toolResult', any>) => ({
+  toolResult: {
+    toolUseId: toolResult.toolUseId,
+    content: toolResult.content.map((toolResultBlock: Record<string, any>) => {
+      if (toolResultBlock.image) {
+        return deserializeImageBlock(toolResultBlock);
+      }
+      if (toolResultBlock.json) {
+        return deserializeJsonBlock(toolResultBlock);
+      }
+      return removeNullsFromBlock(toolResultBlock);
+    }),
+  },
+});
+
+const removeNullsFromBlock = (block: Record<string, any>) =>
+  Object.fromEntries(Object.entries(block).filter(([_, v]) => v !== null));

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

@@ -0,0 +1,68 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { toBase64 } from '@smithy/util-base64';
+import type {
+  ConversationMessageContent,
+  ConversationMessageImageContent,
+  ConversationMessageToolResultContent,
+} from '../../../ai/types/ConversationMessageContent';
+import type { ToolConfiguration } from '../../../ai/types/ToolConfiguration';
+import type { ToolResultJsonContent } from '../../../ai/types/ToolResultContent';
+
+export const serializeAiContext = (aiContext: string | Record<string, any>) =>
+  JSON.stringify(aiContext);
+
+export const serializeContent = (content: ConversationMessageContent[]) =>
+  content.map((block) => {
+    if (block.image) {
+      return serializeImageBlock(block);
+    }
+    if (block.toolResult) {
+      return serializeToolResultBlock(block);
+    }
+    return block;
+  });
+
+export const serializeToolConfiguration = ({ tools }: ToolConfiguration) => ({
+  tools: Object.entries(tools).map(([name, tool]) => ({
+    toolSpec: {
+      name,
+      description: tool.description,
+      inputSchema: {
+        json: JSON.stringify(tool.inputSchema.json),
+      },
+    },
+  })),
+});
+
+const serializeImageBlock = ({ image }: ConversationMessageImageContent) => ({
+  image: {
+    ...image,
+    source: {
+      ...image.source,
+      bytes: toBase64(image.source.bytes),
+    },
+  },
+});
+
+const serializeJsonBlock = ({ json }: ToolResultJsonContent) => ({
+  json: JSON.stringify(json),
+});
+
+const serializeToolResultBlock = ({
+  toolResult,
+}: ConversationMessageToolResultContent) => ({
+  toolResult: {
+    ...toolResult,
+    content: toolResult.content.map((toolResultBlock) => {
+      if (toolResultBlock.image) {
+        return serializeImageBlock(toolResultBlock);
+      }
+      if (toolResultBlock.json) {
+        return serializeJsonBlock(toolResultBlock);
+      }
+      return toolResultBlock;
+    }),
+  },
+});

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

@@ -0,0 +1,53 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { Conversation } from '../../../ai/ConversationType';
+import type {
+  BaseBrowserClient,
+  BaseClient,
+  ClientInternalsGetter,
+  ModelIntrospectionSchema,
+  SchemaModel,
+} from '../../bridge-types';
+import { createListMessagesFunction } from './createListMessagesFunction';
+import { createOnMessageFunction } from './createOnMessageFunction';
+import { createSendMessageFunction } from './createSendMessageFunction';
+
+export const convertItemToConversation = (
+  client: BaseClient,
+  modelIntrospection: ModelIntrospectionSchema,
+  conversationId: string,
+  conversationRouteName: string,
+  conversationMessageModel: SchemaModel,
+  getInternals: ClientInternalsGetter,
+): Conversation => {
+  if (!conversationId) {
+    throw new Error(
+      `An error occurred converting a ${conversationRouteName} conversation: Missing ID`,
+    );
+  }
+  return {
+    id: conversationId,
+    onMessage: createOnMessageFunction(
+      client as BaseBrowserClient,
+      modelIntrospection,
+      conversationId,
+      conversationRouteName,
+      getInternals,
+    ),
+    sendMessage: createSendMessageFunction(
+      client,
+      modelIntrospection,
+      conversationId,
+      conversationRouteName,
+      getInternals,
+    ),
+    listMessages: createListMessagesFunction(
+      client,
+      modelIntrospection,
+      conversationId,
+      conversationMessageModel,
+      getInternals,
+    ),
+  };
+};

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

@@ -0,0 +1,18 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { deserializeContent } from './conversationMessageDeserializers';
+
+export const convertItemToConversationMessage = ({
+  content,
+  createdAt,
+  id,
+  conversationId,
+  role,
+}: any) => ({
+  content: deserializeContent(content),
+  conversationId,
+  createdAt,
+  id,
+  role,
+});

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

@@ -0,0 +1,44 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { ConversationRoute } from '../../../ai/ConversationType';
+import type { SingularReturnValue } from '../../../runtime/client';
+import type {
+  BaseClient,
+  ClientInternalsGetter,
+  ModelIntrospectionSchema,
+  SchemaModel,
+} from '../../bridge-types';
+import { getFactory } from '../operations/get';
+import { convertItemToConversation } from './convertItemToConversation';
+
+export const createCreateConversationFunction =
+  (
+    client: BaseClient,
+    modelIntrospection: ModelIntrospectionSchema,
+    conversationRouteName: string,
+    conversationModel: SchemaModel,
+    conversationMessageModel: SchemaModel,
+    getInternals: ClientInternalsGetter,
+  ): ConversationRoute['create'] =>
+  async () => {
+    const get = getFactory(
+      client,
+      modelIntrospection,
+      conversationModel,
+      'CREATE',
+      getInternals,
+    ) as () => SingularReturnValue<Record<string, any>>;
+    const { data, errors } = await get();
+    return {
+      data: convertItemToConversation(
+        client,
+        modelIntrospection,
+        data?.id,
+        conversationRouteName,
+        conversationMessageModel,
+        getInternals,
+      ),
+      errors,
+    };
+  };

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

@@ -0,0 +1,48 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { ConversationRoute } from '../../../ai/ConversationType';
+import type { SingularReturnValue } from '../../../runtime/client';
+import type {
+  BaseClient,
+  ClientInternalsGetter,
+  ModelIntrospectionSchema,
+  SchemaModel,
+} from '../../bridge-types';
+import { getFactory } from '../operations/get';
+import { convertItemToConversation } from './convertItemToConversation';
+
+export const createGetConversationFunction =
+  (
+    client: BaseClient,
+    modelIntrospection: ModelIntrospectionSchema,
+    conversationRouteName: string,
+    conversationModel: SchemaModel,
+    conversationMessageModel: SchemaModel,
+    getInternals: ClientInternalsGetter,
+  ): ConversationRoute['get'] =>
+  async ({ id }) => {
+    const get = getFactory(
+      client,
+      modelIntrospection,
+      conversationModel,
+      'GET',
+      getInternals,
+    ) as (
+      args?: Record<string, any>,
+    ) => SingularReturnValue<Record<string, any>>;
+    const { data, errors } = await get({ id });
+    return {
+      data: data
+        ? convertItemToConversation(
+            client,
+            modelIntrospection,
+            data.id,
+            conversationRouteName,
+            conversationMessageModel,
+            getInternals,
+          )
+        : data,
+      errors,
+    };
+  };

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

@@ -0,0 +1,46 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { ConversationRoute } from '../../../ai/ConversationType';
+import type { ListReturnValue } from '../../../runtime/client';
+import type {
+  BaseClient,
+  ClientInternalsGetter,
+  ModelIntrospectionSchema,
+  SchemaModel,
+} from '../../bridge-types';
+import { listFactory } from '../operations/list';
+import { convertItemToConversation } from './convertItemToConversation';
+
+export const createListConversationsFunction =
+  (
+    client: BaseClient,
+    modelIntrospection: ModelIntrospectionSchema,
+    conversationRouteName: string,
+    conversationModel: SchemaModel,
+    conversationMessageModel: SchemaModel,
+    getInternals: ClientInternalsGetter,
+  ): ConversationRoute['list'] =>
+  async (input) => {
+    const list = listFactory(
+      client,
+      modelIntrospection,
+      conversationModel,
+      getInternals,
+    ) as (args?: Record<string, any>) => ListReturnValue<Record<string, any>>;
+    const { data, nextToken, errors } = await list(input);
+    return {
+      data: data.map((datum: Record<string, any>) => {
+        return convertItemToConversation(
+          client,
+          modelIntrospection,
+          datum.id,
+          conversationRouteName,
+          conversationMessageModel,
+          getInternals,
+        );
+      }),
+      nextToken,
+      errors,
+    };
+  };

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

@@ -0,0 +1,42 @@
+// 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 { ListReturnValue } from '../../../runtime/client';
+import type {
+  BaseClient,
+  ClientInternalsGetter,
+  ModelIntrospectionSchema,
+  SchemaModel,
+} from '../../bridge-types';
+import { listFactory } from '../operations/list';
+import { convertItemToConversationMessage } from './convertItemToConversationMessage';
+
+export const createListMessagesFunction =
+  (
+    client: BaseClient,
+    modelIntrospection: ModelIntrospectionSchema,
+    conversationId: string,
+    conversationMessageModel: SchemaModel,
+    getInternals: ClientInternalsGetter,
+  ): Conversation['listMessages'] =>
+  async (input) => {
+    const list = listFactory(
+      client,
+      modelIntrospection,
+      conversationMessageModel,
+      getInternals,
+    ) as (args?: Record<string, any>) => ListReturnValue<ConversationMessage>;
+    const { data, nextToken, errors } = await list({
+      ...input,
+      filter: { conversationId: { eq: conversationId } },
+    });
+    return {
+      data: data.map((item: any) => convertItemToConversationMessage(item)),
+      nextToken,
+      errors,
+    };
+  };

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

@@ -0,0 +1,41 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { Observable, Subscription } from 'rxjs';
+import type { Conversation } from '../../../ai/ConversationType';
+import type {
+  BaseClient,
+  ClientInternalsGetter,
+  ModelIntrospectionSchema,
+} from '../../bridge-types';
+import { customOpFactory } from '../operations/custom';
+import { convertItemToConversationMessage } from './convertItemToConversationMessage';
+
+export const createOnMessageFunction =
+  (
+    client: BaseClient,
+    modelIntrospection: ModelIntrospectionSchema,
+    conversationId: string,
+    conversationRouteName: string,
+    getInternals: ClientInternalsGetter,
+  ): Conversation['onMessage'] =>
+  (handler): Subscription => {
+    const { conversations } = modelIntrospection;
+    // Safe guard for standalone function. When called as part of client generation, this should never be falsy.
+    if (!conversations) {
+      return {} as Subscription;
+    }
+    const subscribeSchema =
+      conversations[conversationRouteName].message.subscribe;
+    const subscribeOperation = customOpFactory(
+      client,
+      modelIntrospection,
+      'subscription',
+      subscribeSchema,
+      false,
+      getInternals,
+    ) as (args?: Record<string, any>) => Observable<any>;
+    return subscribeOperation({ conversationId }).subscribe((data) => {
+      handler(convertItemToConversationMessage(data));
+    });
+  };