@aws-amplify/data-schema 1.4.0 → 1.5.0

package/src/ai/ConversationSchemaTypes.ts

@@ -0,0 +1,253 @@
+import { InternalRef } from '../RefType';
+import { capitalize } from '../runtime/utils';
+import type {
+  InternalConversationType,
+  ToolDefinition,
+} from './ConversationType';
+
+export const createConversationField = (
+  typeDef: InternalConversationType,
+  typeName: string,
+): string => {
+  const { aiModel, systemPrompt, handler, tools } = typeDef;
+
+  const args: Record<string, string> = {
+    aiModel: aiModel.resourcePath,
+    systemPrompt,
+  };
+
+  if (handler) {
+    if (typeof handler === 'string') {
+      args['functionName'] = handler;
+    } else if (typeof handler.getInstance === 'function') {
+      args['functionName'] = `Fn${capitalize(typeName)}`;
+    }
+  }
+
+  const argsString = Object.entries(args)
+    .map(([key, value]) => `${key}: "${value}"`)
+    .join(', ');
+
+  const toolsString = tools?.length
+    ? `, tools: [${getConversationToolsString(tools)}]`
+    : '';
+
+  const conversationDirective = `@conversation(${argsString}${toolsString})`;
+
+  return `${typeName}(conversationId: ID!, content: [ContentBlockInput], aiContext: AWSJSON, toolConfiguration: ToolConfigurationInput): ConversationMessage ${conversationDirective} @aws_cognito_user_pools`;
+};
+
+const isRef = (query: unknown): query is { data: InternalRef['data'] } =>
+  (query as any)?.data?.type === 'ref';
+
+const getConversationToolsString = (tools: ToolDefinition[]) =>
+  tools
+    .map((tool) => {
+      const { query, description } = tool;
+      if (!isRef(query)) {
+        throw new Error(`Unexpected query was found in tool ${tool}.`);
+      }
+      // TODO: add validation for query / auth (cup) / etc
+      const queryName = query.data.link;
+      return `{ name: "${queryName}", description: "${description}" }`;
+    })
+    .join(', ');
+
+const ConversationParticipantRole = `enum ConversationParticipantRole {
+  user
+  assistant
+}`;
+
+const ConversationMessage = `interface ConversationMessage {
+  id: ID!
+  conversationId: ID!
+  role: ConversationParticipantRole
+  content: [ContentBlock]
+  aiContext: AWSJSON
+  toolConfiguration: ToolConfiguration
+  createdAt: AWSDateTime
+  updatedAt: AWSDateTime
+  owner: String
+}`;
+
+const DocumentBlockSourceInput = `input DocumentBlockSourceInput {
+  bytes: String
+}`;
+
+const DocumentBlockInput = `input DocumentBlockInput {
+  format: String!
+  name: String!
+  source: DocumentBlockSourceInput!
+}`;
+
+const ImageBlockSourceInput = `input ImageBlockSourceInput {
+  bytes: String
+}`;
+
+const ImageBlockInput = `input ImageBlockInput {
+  format: String!
+  source: ImageBlockSourceInput!
+}`;
+
+const ToolResultContentBlockInput = `input ToolResultContentBlockInput {
+  document: DocumentBlockInput
+  image: ImageBlockInput
+  json: AWSJSON
+  text: String
+}`;
+
+const ToolResultBlockInput = `input ToolResultBlockInput {
+  content: [ToolResultContentBlockInput!]!
+  toolUseId: String!
+  status: String
+}`;
+
+const DocumentBlockSource = `type DocumentBlockSource {
+  bytes: String
+}
+`;
+const DocumentBlock = `type DocumentBlock {
+  format: String!
+  name: String!
+  source: DocumentBlockSource!
+}`;
+
+const ImageBlock = `type ImageBlock {
+  format: String!
+  source: ImageBlockSource!
+}`;
+
+const ImageBlockSource = `type ImageBlockSource {
+  bytes: String
+}`;
+
+const ToolUseBlockInput = `input ToolUseBlockInput {
+  toolUseId: String!
+  name: String!
+  input: AWSJSON!
+}`;
+
+const ToolUseBlock = `type ToolUseBlock {
+  toolUseId: String!
+  name: String!
+  input: AWSJSON!
+}`;
+
+const ToolResultContentBlock = `type ToolResultContentBlock {
+  document: DocumentBlock
+  image: ImageBlock
+  json: AWSJSON
+  text: String
+}`;
+
+const ToolResultBlock = `type ToolResultBlock {
+  content: [ToolResultContentBlock!]!
+  toolUseId: String!
+  status: String
+}`;
+
+const ContentBlockText = `type ContentBlockText {
+  text: String
+}`;
+
+const ContentBlockImage = `type ContentBlockImage {
+  image: ImageBlock
+}`;
+
+const ContentBlockDocument = `type ContentBlockDocument {
+  document: DocumentBlock
+}`;
+
+const ContentBlockToolUse = `type ContentBlockToolUse {
+  toolUse: ToolUseBlock
+}`;
+
+const ContentBlockToolResult = `type ContentBlockToolResult {
+  toolResult: ToolResultBlock
+}`;
+
+const ContentBlockInput = `input ContentBlockInput {
+  text: String
+  document: DocumentBlockInput
+  image: ImageBlockInput
+  toolResult: ToolResultBlockInput
+  toolUse: ToolUseBlockInput
+}`;
+
+const ContentBlock = `type ContentBlock {
+  text: String
+  document: DocumentBlock
+  image: ImageBlock
+  toolResult: ToolResultBlock
+  toolUse: ToolUseBlock
+}`;
+
+const ToolConfigurationInput = `input ToolConfigurationInput {
+  tools: [ToolInput]
+}`;
+
+const ToolInput = `input ToolInput {
+  toolSpec: ToolSpecificationInput
+}`;
+
+const ToolSpecificationInput = `input ToolSpecificationInput {
+  name: String!
+  description: String
+  inputSchema: ToolInputSchemaInput!
+}`;
+
+const ToolInputSchemaInput = `input ToolInputSchemaInput {
+  json: AWSJSON
+}`;
+
+const ToolConfiguration = `type ToolConfiguration {
+  tools: [Tool]
+}`;
+
+const Tool = `type Tool {
+  toolSpec: ToolSpecification
+}`;
+
+const ToolSpecification = `type ToolSpecification {
+  name: String!
+  description: String
+  inputSchema: ToolInputSchema!
+}`;
+
+const ToolInputSchema = `type ToolInputSchema {
+  json: AWSJSON
+}`;
+
+export const conversationTypes: string[] = [
+  ConversationParticipantRole,
+  ConversationMessage,
+  DocumentBlockSourceInput,
+  DocumentBlockInput,
+  ImageBlockSourceInput,
+  ImageBlockInput,
+  ToolUseBlockInput,
+  ToolResultContentBlockInput,
+  ToolResultBlockInput,
+  DocumentBlockSource,
+  DocumentBlock,
+  ImageBlock,
+  ImageBlockSource,
+  ToolUseBlock,
+  ToolResultContentBlock,
+  ToolResultBlock,
+  ContentBlockText,
+  ContentBlockImage,
+  ContentBlockDocument,
+  ContentBlockToolUse,
+  ContentBlockToolResult,
+  ContentBlockInput,
+  ContentBlock,
+  ToolConfigurationInput,
+  ToolInput,
+  ToolSpecificationInput,
+  ToolInputSchemaInput,
+  ToolConfiguration,
+  Tool,
+  ToolSpecification,
+  ToolInputSchema,
+];

package/src/ai/ConversationType.ts

@@ -0,0 +1,135 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { Subscription } from 'rxjs';
+import type { DefineFunction } from '@aws-amplify/data-schema-types';
+import type { ListReturnValue, SingularReturnValue } from '../runtime/client';
+import { type Brand, brand } from '../util';
+import type { RefType } from '../RefType';
+import { AiModel, InferenceConfiguration } from './ModelType';
+import {
+  ConversationMessageContent,
+  ConversationSendMessageInputContent,
+} from './types/ConversationMessageContent';
+import { ToolConfiguration } from './types/ToolConfiguration';
+
+export const brandName = 'conversationCustomOperation';
+
+// conversation message types
+export interface ConversationMessage {
+  content: ConversationMessageContent[];
+  conversationId: string;
+  createdAt: string;
+  id: string;
+  role: 'user' | 'assistant';
+}
+
+// conversation route types
+interface ConversationRouteGetInput {
+  id: string;
+}
+
+interface ConversationRouteListInput {
+  limit?: number;
+  nextToken?: string | null;
+}
+
+export interface ConversationRoute {
+  /**
+   * @experimental
+   *
+   * Creates a {@link Conversation} from the current conversation route.
+   */
+  create: () => SingularReturnValue<Conversation>;
+  /**
+   * @experimental
+   *
+   * Gets an existing {@link Conversation} based on ID.
+   */
+  get: (input: ConversationRouteGetInput) => SingularReturnValue<Conversation>;
+  /**
+   * @experimental
+   *
+   * Lists all existing {@link Conversation}s on the current conversation route.
+   */
+  list: (input?: ConversationRouteListInput) => ListReturnValue<Conversation>;
+}
+
+// conversation types
+interface ConversationSendMessageInput {
+  content: ConversationSendMessageInputContent[];
+  aiContext?: string | Record<string, any>;
+  toolConfiguration?: ToolConfiguration;
+}
+
+interface ConversationListMessagesInput {
+  limit?: number;
+  nextToken?: string | null;
+}
+
+type ConversationOnMessageHandler = (message: ConversationMessage) => void;
+
+export interface Conversation {
+  id: string;
+  /**
+   * @experimental
+   *
+   * Sends a message to the current conversation.
+   */
+  sendMessage: (
+    input: ConversationSendMessageInput,
+  ) => SingularReturnValue<ConversationMessage>;
+  /**
+   * @experimental
+   *
+   * Lists all existing messages for the current conversation.
+   */
+  listMessages: (
+    input?: ConversationListMessagesInput,
+  ) => ListReturnValue<ConversationMessage>;
+  /**
+   * @experimental
+   *
+   * Subscribes to new messages on the current conversation.
+   */
+  onMessage: (handler: ConversationOnMessageHandler) => Subscription;
+}
+
+// schema definition input
+export interface ToolDefinition {
+  query: RefType<any>;
+  description: string;
+}
+
+export interface ConversationInput {
+  aiModel: AiModel;
+  systemPrompt: string;
+  inferenceConfiguration?: InferenceConfiguration;
+  tools?: ToolDefinition[];
+  handler?: DefineFunction | string;
+}
+
+export interface InternalConversationType
+  extends ConversationType,
+    ConversationInput {}
+
+export interface ConversationType extends Brand<typeof brandName> {}
+
+function _conversation(input: ConversationInput): ConversationType {
+  return { ...brand(brandName), ...input };
+}
+
+/**
+ * @experimental
+ *
+ * Define an AI conversation route which enables multi-turn conversation APIs for interacting with specified AI model.
+ * @example
+ * realtorChat: a.conversation({
+ *   aiModel: { resourcePath },
+ *   systemPrompt: 'You are a helpful real estate assistant',
+ * })
+ * @returns a conversation route definition
+ */
+export function conversation(input: ConversationInput): ConversationType {
+  return _conversation(input);
+}

package/src/ai/ModelType.ts

@@ -0,0 +1,43 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+const supportedModelsLookup = {
+  // Anthropic models
+  'Claude 3 Haiku': 'anthropic.claude-3-haiku-20240307-v1:0',
+  'Claude 3 Opus': 'anthropic.claude-3-opus-20240229-v1:0',
+  'Claude 3 Sonnet': 'anthropic.claude-3-sonnet-20240229-v1:0',
+  'Claude 3.5 Sonnet': 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+  // Cohere models
+  'Cohere Command R': 'cohere.command-r-v1:0',
+  'Cohere Command R+': 'cohere.command-r-plus-v1:0',
+  // Meta models
+  'Llama 3.1 8B Instruct': 'meta.llama3-1-8b-instruct-v1:0',
+  'Llama 3.1 70B Instruct': 'meta.llama3-1-70b-instruct-v1:0',
+  'Llama 3.1 405B Instruct': 'meta.llama3-1-405b-instruct-v1:0',
+  // Mistral AI models
+  'Mistral Large': 'mistral.mistral-large-2402-v1:0',
+  'Mistral Large 2': 'mistral.mistral-large-2407-v1:0',
+  'Mistral Small': 'mistral.mistral-small-2402-v1:0',
+} as const;
+
+export interface AiModel {
+  resourcePath: string;
+}
+
+export interface InferenceConfiguration {
+  topP?: number;
+  temperature?: number;
+  maxTokens?: number;
+}
+
+/**
+ * @experimental
+ *
+ * Bedrock models currently supporting Converse API and Tool use
+ * @see {@link https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html#conversation-inference-supported-models-features}
+ */
+export function model(modelName: keyof typeof supportedModelsLookup): AiModel {
+  return {
+    resourcePath: supportedModelsLookup[modelName],
+  };
+}

package/src/ai/types/ConversationMessageContent.ts

@@ -0,0 +1,43 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { ImageBlock, ToolResultBlock, ToolUseBlock } from './contentBlocks';
+
+export interface ConversationMessageTextContent {
+  text: string;
+  image?: never;
+  toolUse?: never;
+  toolResult?: never;
+}
+
+export interface ConversationMessageImageContent {
+  text?: never;
+  image: ImageBlock;
+  toolUse?: never;
+  toolResult?: never;
+}
+
+export interface ConversationMessageToolUseContent {
+  text?: never;
+  image?: never;
+  toolUse: ToolUseBlock;
+  toolResult?: never;
+}
+
+export interface ConversationMessageToolResultContent {
+  text?: never;
+  image?: never;
+  toolUse?: never;
+  toolResult: ToolResultBlock;
+}
+
+export type ConversationMessageContent =
+  | ConversationMessageTextContent
+  | ConversationMessageImageContent
+  | ConversationMessageToolUseContent
+  | ConversationMessageToolResultContent;
+
+export type ConversationSendMessageInputContent =
+  | Omit<ConversationMessageTextContent, 'toolUse'>
+  | Omit<ConversationMessageImageContent, 'toolUse'>
+  | Omit<ConversationMessageToolResultContent, 'toolUse'>;

package/src/ai/types/ToolConfiguration.ts

@@ -0,0 +1,24 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { JSONSchema4 } from 'json-schema';
+
+interface ObjectTypedJSONSchema4 extends JSONSchema4 {
+  type: 'object';
+}
+
+interface ToolJsonInputSchema {
+  /**
+   * The schema for the tool. The top level schema type must be object.
+   */
+  json: ObjectTypedJSONSchema4;
+}
+
+export interface Tool {
+  inputSchema: ToolJsonInputSchema;
+  description?: string;
+}
+
+export interface ToolConfiguration {
+  tools: Record<string, Tool>;
+}

package/src/ai/types/ToolResultContent.ts

@@ -0,0 +1,28 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { DocumentType } from '../../runtime/bridge-types';
+import { ImageBlock } from './contentBlocks';
+
+export interface ToolResultTextContent {
+  text: string;
+  image?: never;
+  json?: never;
+}
+
+export interface ToolResultImageContent {
+  text?: never;
+  image: ImageBlock;
+  json?: never;
+}
+
+export interface ToolResultJsonContent {
+  text?: never;
+  image?: never;
+  json: DocumentType;
+}
+
+export type ToolResultContent =
+  | ToolResultTextContent
+  | ToolResultImageContent
+  | ToolResultJsonContent;

package/src/ai/types/contentBlocks.ts

@@ -0,0 +1,26 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import type { DocumentType } from '../../runtime/bridge-types';
+import type { ToolResultContent } from './ToolResultContent';
+
+interface BytesImageSource {
+  bytes: Uint8Array;
+}
+
+// common content blocks
+export interface ImageBlock {
+  format: 'gif' | 'jpeg' | 'png' | 'webp';
+  source: BytesImageSource;
+}
+
+export interface ToolUseBlock {
+  toolUseId: string;
+  name: string;
+  input: DocumentType;
+}
+
+export interface ToolResultBlock {
+  toolUseId: string;
+  content: ToolResultContent[];
+}

package/src/runtime/addSchemaToClient.ts

@@ -3,10 +3,12 @@
 
 import { ClientExtensions } from './client';
 import {
+  generateConversationsProperty,
   generateCustomQueriesProperty,
   generateCustomMutationsProperty,
   generateCustomSubscriptionsProperty,
   generateEnumsProperty,
+  generateGenerationsProperty,
   generateModelsProperty,
   upgradeClientCancellation,
 } from './internals';
@@ -43,5 +45,15 @@ export function addSchemaToClient<T extends Record<any, any> = never>(
     apiGraphqlConfig,
     getInternals,
   );
+  (client as any).conversations = generateConversationsProperty(
+    client as any,
+    apiGraphqlConfig,
+    getInternals,
+  );
+  (client as any).generations = generateGenerationsProperty(
+    client as any,
+    apiGraphqlConfig,
+    getInternals,
+  );
   return client as any;
 }

package/src/runtime/bridge-types.ts

@@ -87,6 +87,8 @@ export interface ModelIntrospectionSchema {
   queries?: CustomOperations;
   mutations?: CustomOperations;
   subscriptions?: CustomOperations;
+  conversations?: SchemaConversationRoutes;
+  generations?: SchemaGenerationRoutes;
 }
 
 /**
@@ -96,6 +98,27 @@ export type SchemaModels = Record<string, SchemaModel>;
 export type SchemaNonModels = Record<string, SchemaNonModel>;
 export type SchemaEnums = Record<string, SchemaEnum>;
 export type CustomOperations = Record<string, CustomOperation>;
+export type SchemaConversationRoutes = Record<string, SchemaConversationRoute>;
+export type SchemaGenerationRoutes = Record<string, CustomOperation>;
+
+export type SchemaConversationRoute = {
+  name: string;
+  models: SchemaModels;
+  nonModels: SchemaNonModels;
+  enums: SchemaEnums;
+  conversation: SchemaConversation;
+  message: SchemaConversationMessage;
+};
+
+export type SchemaConversation = {
+  modelName: string;
+};
+
+export type SchemaConversationMessage = {
+  modelName: string;
+  subscribe: CustomOperation;
+  send: CustomOperation;
+};
 
 export interface SchemaModel {
   name: string;
@@ -140,7 +163,7 @@ export type CustomOperationArguments = Record<string, CustomOperationArgument>;
 
 export interface CustomOperationArgument {
   name: string;
-  type: FieldType;
+  type: InputFieldType;
   isArray: boolean;
   isRequired: boolean;
   isArrayNullable?: boolean;
@@ -185,7 +208,15 @@ export interface NonModelFieldType {
   nonModel: string;
 }
 
-export type FieldType =
+export interface EnumType {
+  enum: string;
+}
+
+export interface InputType {
+  input: string;
+}
+
+type ScalarType =
   | 'ID'
   | 'String'
   | 'Int'
@@ -199,10 +230,16 @@ export type FieldType =
   | 'AWSIPAddress'
   | 'Boolean'
   | 'AWSJSON'
-  | 'AWSPhone'
-  | { enum: string }
+  | 'AWSPhone';
+
+export type FieldType =
+  | ScalarType
+  | EnumType
   | ModelFieldType
   | NonModelFieldType;
+
+export type InputFieldType = ScalarType | EnumType | InputType;
+
 export type FieldAttribute = ModelAttribute;
 
 /**