AI streaming (#5920)

.changeset/curvy-tables-work.md

@@ -0,0 +1,7 @@
+---
+"@aws-amplify/ui-react-ai": minor
+---
+
+AI streaming
+
+This update brings streaming responses to the `useAIConversation()` hook and `<AIConversation />` component. No changes need to be made in your code to get streaming responses now. 

examples/next/pages/ui/components/ai/ai-conversation-response-components/[id].page.tsx

@@ -0,0 +1,113 @@
+import * as React from 'react';
+import { Amplify } from 'aws-amplify';
+import { createAIHooks, AIConversation } from '@aws-amplify/ui-react-ai';
+import { generateClient } from 'aws-amplify/api';
+import ReactMarkdown from 'react-markdown';
+import '@aws-amplify/ui-react/styles.css';
+import { GlobalStyle } from '@aws-amplify/ui-react/server';
+
+import outputs from './amplify_outputs';
+import type { Schema } from '@environments/ai/gen2/amplify/data/resource';
+import { Authenticator, Button, Card } from '@aws-amplify/ui-react';
+import { useRouter } from 'next/router';
+
+const client = generateClient<Schema>({ authMode: 'userPool' });
+const { useAIConversation } = createAIHooks(client);
+
+Amplify.configure(outputs);
+
+const AIContext = React.createContext<any>(undefined);
+
+const responseComponents = {
+  WeatherCard: {
+    description: 'Used to display the weather in a city',
+    component: ({ city }) => {
+      const aiContext = React.useContext(AIContext);
+      aiContext.setAIContext({
+        temperature: 72,
+        city,
+        weather: 'sunny',
+      });
+      return (
+        <Card variation="outlined">
+          <h3>Weather Card {city}</h3>
+        </Card>
+      );
+    },
+    props: {
+      city: {
+        type: 'string',
+        description: 'The city to get the weather for',
+      },
+    },
+  },
+} as const;
+
+function Chat({ id }: { id?: string }) {
+  const aiContext = React.useContext(AIContext);
+  const [
+    {
+      data: { messages },
+      isLoading,
+    },
+    sendMessage,
+  ] = useAIConversation('pirateChat', {
+    id,
+  });
+
+  console.log(aiContext);
+
+  return (
+    <AIConversation
+      messages={messages}
+      isLoading={isLoading}
+      handleSendMessage={(message) => {
+        sendMessage({
+          ...message,
+          aiContext: aiContext.aiContext,
+        });
+      }}
+      messageRenderer={{
+        text: ({ text }) => <ReactMarkdown>{text}</ReactMarkdown>,
+      }}
+      responseComponents={responseComponents}
+    />
+  );
+}
+
+export default function Example() {
+  const router = useRouter();
+  const [shown, setShown] = React.useState(true);
+  const [aiContext, setAIContext] = React.useState<any>(undefined);
+
+  return (
+    <Authenticator>
+      <Button
+        onClick={() => {
+          setShown(!shown);
+        }}
+      >
+        Toggle
+      </Button>
+      <AIContext.Provider value={{ aiContext, setAIContext }}>
+        {shown ? (
+          <Card variation="outlined" height="400px" margin="large">
+            <Chat id={`${router.query.id}`} />
+          </Card>
+        ) : null}
+      </AIContext.Provider>
+      <GlobalStyle
+        styles={{
+          code: {
+            backgroundColor: 'var(--amplify-colors-background-tertiary)',
+          },
+          pre: {
+            backgroundColor: 'var(--amplify-colors-background-tertiary)',
+            display: 'block',
+            padding: '1rem',
+          },
+        }}
+      />
+    </Authenticator>
+  );
+}

examples/next/pages/ui/components/ai/ai-conversation/[id].page.tsx

@@ -0,0 +1,64 @@
+import * as React from 'react';
+import { Amplify } from 'aws-amplify';
+import { createAIHooks, AIConversation } from '@aws-amplify/ui-react-ai';
+import { generateClient } from 'aws-amplify/api';
+import '@aws-amplify/ui-react/styles.css';
+
+import outputs from './amplify_outputs';
+import type { Schema } from '@environments/ai/gen2/amplify/data/resource';
+import { Authenticator, Button, Card } from '@aws-amplify/ui-react';
+import { useRouter } from 'next/router';
+
+const client = generateClient<Schema>({ authMode: 'userPool' });
+const { useAIConversation } = createAIHooks(client);
+
+Amplify.configure(outputs);
+
+const onMessage = (message) => {
+  console.log('onmessage', message);
+};
+
+function Chat({ id }: { id?: string }) {
+  const [
+    {
+      data: { messages },
+      isLoading,
+    },
+    sendMessage,
+  ] = useAIConversation('pirateChat', {
+    id,
+    onMessage,
+  });
+
+  return (
+    <AIConversation
+      messages={messages}
+      isLoading={isLoading}
+      handleSendMessage={sendMessage}
+    />
+  );
+}
+
+export default function Example() {
+  const router = useRouter();
+  const [shown, setShown] = React.useState(true);
+
+  if (router.query.id) {
+    return (
+      <Authenticator>
+        <Button
+          onClick={() => {
+            setShown(!shown);
+          }}
+        >
+          Toggle
+        </Button>
+        {shown ? (
+          <Card variation="outlined" height="400px" margin="large">
+            <Chat id={`${router.query.id}`} />
+          </Card>
+        ) : null}
+      </Authenticator>
+    );
+  }
+}

examples/next/pages/ui/components/ai/ai-conversation/index.page.tsx

@@ -1,63 +1,108 @@
+import * as React from 'react';
 import { Amplify } from 'aws-amplify';
 import { createAIHooks, AIConversation } from '@aws-amplify/ui-react-ai';
 import { generateClient } from 'aws-amplify/api';
 import '@aws-amplify/ui-react/styles.css';
 
 import outputs from './amplify_outputs';
 import type { Schema } from '@environments/ai/gen2/amplify/data/resource';
-import { Authenticator, Card } from '@aws-amplify/ui-react';
+import {
+  Authenticator,
+  Button,
+  Card,
+  Flex,
+  Heading,
+  View,
+} from '@aws-amplify/ui-react';
+import { useRouter } from 'next/router';
 
 const client = generateClient<Schema>({ authMode: 'userPool' });
 const { useAIConversation } = createAIHooks(client);
 
 Amplify.configure(outputs);
 
-const formatDate = (date: Date): string =>
-  `Argh the time be round ${date.toLocaleTimeString('en-US', {
-    hour: 'numeric',
-    minute: 'numeric',
-    hour12: true,
-  })}`;
-
 function Chat() {
   const [
     {
       data: { messages },
       isLoading,
     },
     sendMessage,
-  ] = useAIConversation('pirateChat');
+  ] = useAIConversation('pirateChat', { onMessage, onInitialize });
+
+  return (
+    <AIConversation
+      messages={messages}
+      isLoading={isLoading}
+      handleSendMessage={sendMessage}
+    />
+  );
+}
+
+const onInitialize = (conversation) => {
+  console.log(conversation);
+};
+
+const onMessage = (conversation) => {
+  console.log('on message');
+  console.log(conversation);
+};
+
+function SyncedChats() {
+  const [
+    {
+      data: { messages },
+      isLoading,
+    },
+    handleSendMessage,
+  ] = useAIConversation('pirateChat', { onInitialize, onMessage });
+
+  const props = {
+    isLoading,
+    handleSendMessage,
+    messages,
+  };
 
   return (
-    <Card variation="outlined" width="50%" height="300px" margin="0 auto">
-      <AIConversation
-        displayText={{ getMessageTimestampText: formatDate }}
-        messages={messages}
-        handleSendMessage={sendMessage}
-        isLoading={isLoading}
-        allowAttachments
-        suggestedPrompts={[
-          {
-            inputText: 'hello',
-            component: 'hello',
-          },
-          {
-            inputText: 'how are you?',
-            component: 'how are you?',
-          },
-        ]}
-        variant="bubble"
-      />
-    </Card>
+    <Flex direction="row">
+      <Card flex="1" variation="outlined" height="400px" margin="large">
+        <AIConversation {...props} />
+      </Card>
+      <Card flex="1" variation="outlined" height="400px" margin="large">
+        <AIConversation {...props} />
+      </Card>
+    </Flex>
   );
 }
 
 export default function Example() {
+  const router = useRouter();
+  const handleCreateChat = async () => {
+    const { data } = await client.conversations.pirateChat.create();
+    if (data.id) {
+      router.push(`/ui/components/ai/ai-conversation/${data.id}`);
+    }
+  };
   return (
     <Authenticator>
-      {({ user, signOut }) => {
-        return <Chat />;
-      }}
+      <Flex direction="column">
+        <View>
+          <Heading level={2}>Separate chats</Heading>
+          <Flex direction="row">
+            <Card flex="1" variation="outlined" height="400px" margin="large">
+              <Chat />
+            </Card>
+            <Card flex="1" variation="outlined" height="400px" margin="large">
+              <Chat />
+            </Card>
+          </Flex>
+        </View>
+        <View>
+          <Heading level={2}>Synced chats</Heading>
+          <SyncedChats />
+        </View>
+      </Flex>
+      <Button onClick={handleCreateChat}>Create chat</Button>
     </Authenticator>
   );
 }

examples/next/pages/ui/components/ai/ai-conversation/new.page.tsx

@@ -0,0 +1,51 @@
+import * as React from 'react';
+import { Amplify } from 'aws-amplify';
+import { createAIHooks, AIConversation } from '@aws-amplify/ui-react-ai';
+import { generateClient } from 'aws-amplify/api';
+import '@aws-amplify/ui-react/styles.css';
+
+import outputs from './amplify_outputs';
+import type { Schema } from '@environments/ai/gen2/amplify/data/resource';
+import { Authenticator, Card, Flex } from '@aws-amplify/ui-react';
+import { useRouter } from 'next/router';
+
+const client = generateClient<Schema>({ authMode: 'userPool' });
+const { useAIConversation } = createAIHooks(client);
+
+Amplify.configure(outputs);
+
+function Chat() {
+  const router = useRouter();
+
+  const [
+    {
+      data: { messages },
+      isLoading,
+    },
+    sendMessage,
+  ] = useAIConversation('pirateChat', {
+    onInitialize(conversation) {
+      router.replace(`/ui/components/ai/ai-conversation/${conversation.id}`);
+    },
+  });
+
+  return (
+    <AIConversation
+      messages={messages}
+      isLoading={isLoading}
+      handleSendMessage={sendMessage}
+    />
+  );
+}
+
+export default function Example() {
+  return (
+    <Authenticator>
+      <Flex direction="row">
+        <Card flex="1" variation="outlined" height="400px" margin="large">
+          <Chat />
+        </Card>
+      </Flex>
+    </Authenticator>
+  );
+}