AI Components

import {useState, useRef} from 'react';
import {
  AttachFileMenuItem,
  CommandMenuItem,
  InsertTokenMenuItem,
  InsertMenuButton,
  PromptField,
  Attachment,
  AttachmentPreview,
  PromptFieldAttachment,
  PromptFieldAttachmentList,
  PromptFieldSubmitButton,
  PromptFieldToolbar,
  PromptFieldValue,
  PromptFieldVoiceButton,
  PromptToken,
  PromptTokenField
} from '@react-spectrum/ai';
import {type FocusableRefValue} from '@react-types/shared';
import {getIcon, slashCommands, objects, renderCompletions, suggestions, type UploadState} from './ai-component-helpers/promptfield';
import {style} from '@react-spectrum/s2/style' with {type: 'macro'};
import {Collection, SubmenuTrigger, Menu, MenuItem, MenuSection, Header, Heading, Text} from '@react-spectrum/s2';
import Data from '@react-spectrum/s2/icons/Data';
import Plugin from '@react-spectrum/s2/icons/Plugin';
import Prompt from '@react-spectrum/s2/icons/Prompt';
import {VirtualizedStreamingChat} from './ai-component-helpers/chat';

function Example() {
  let [value, setValue] = useState<PromptFieldValue>(() => new PromptFieldValue([]));
  let promptFieldRef = useRef<FocusableRefValue<HTMLDivElement>>(null);
  let [attachments, setAttachments] = useState<PromptFieldAttachment[]>([]);
  let [attachmentState, setAttachmentState] = useState<Map<string, UploadState>>(new Map());

  let mockUpload = async (id: string) => {
    await new Promise(resolve => setTimeout(resolve, Math.random() * 30));
    setAttachmentState(prev => {
      let item = prev.get(id);
      if (!item || item.status === 'completed') {
        return prev;
      }
      let newState = new Map(prev);
      let progress = (item.progress ?? 0) + 1;
      if (progress >= 100) {
        newState.set(id, {status: 'completed'});
      } else {
        newState.set(id, {status: 'uploading', progress});
        mockUpload(id);
      }
      return newState;
    });
  };

  let clearPrompt = () => {
    setValue(new PromptFieldValue([]));
    setAttachments([]);
    alert('Conversation cleared');
  };

  let compactPrompt = () => {
    alert('Conversation compacted');
  };

  return (
    <div className={style({height: 700, width: 'full'})}>
      <VirtualizedStreamingChat
        suggestions={suggestions}
        onSelectSuggestion={value => {
          setValue(value as PromptFieldValue);
          promptFieldRef.current?.focus();
        }}>
        {(onSend, isGenerating) => (
          <PromptField
            ref={promptFieldRef}
            value={value}
            onChange={setValue}
            attachments={attachments}
            onAttachmentsChange={setAttachments}
            isGenerating={isGenerating}
            onSubmit={prompt => {
              onSend(prompt);
              setValue(new PromptFieldValue([]));
              setAttachments([]);
              setAttachmentState(new Map());
            }}
            acceptedAttachmentTypes={['*/*']}
            onAddAttachments={newAttachments => {
              setAttachmentState(prev => {
                let newState = new Map(prev);
                newAttachments.forEach(attachment => {
                  newState.set(attachment.id, {status: 'uploading', progress: 0});
                  mockUpload(attachment.id);
                });
                return newState;
              });
            }}
            onRemoveAttachments={removedAttachments => {
              setAttachmentState(prev => {
                let newState = new Map(prev);
                removedAttachments.forEach(attachment => {
                  newState.delete(attachment.id);
                });
                return newState;
              });
            }}>
            <PromptFieldAttachmentList dependencies={[attachmentState]}>
              {attachment => {
                let state = attachmentState.get(attachment.id);
                return (
                  <Attachment uploadProgress={state?.status === 'uploading' ? state?.progress : undefined}>
                    <AttachmentPreview mimeType={attachment.file.type} src={attachment.image} />
                  </Attachment>
                );
              }}
            </PromptFieldAttachmentList>
            <PromptTokenField
              completionTrigger={/(?<=^|\s)[@/]/}
              renderCompletions={(filterValue, valueType) => {
                return renderCompletions(filterValue, {valueType, onClear: clearPrompt, onCompact: compactPrompt});
              }}>
              {token => (
                <PromptToken token={token}>
                  {getIcon(token)}
                  {token.text}
                </PromptToken>
              )}
            </PromptTokenField>
            <PromptFieldToolbar>
              <div className={style({display: 'flex', gap: 8, alignItems: 'center'})}>
                <InsertMenuButton>
                  <AttachFileMenuItem />
                  <SubmenuTrigger>
                    <MenuItem>
                      <Prompt />
                      <Text>Commands</Text>
                    </MenuItem>
                    <Menu items={slashCommands.filter(item => item.kind === 'command')}>
                      {item => (
                        <CommandMenuItem
                          id={item.command}
                          onAction={item.command === '/clear' ? clearPrompt : compactPrompt}>
                          <Prompt />
                          <Text slot="label">{item.command}</Text>
                          <Text slot="description">{item.description}</Text>
                        </CommandMenuItem>
                      )}
                    </Menu>
                  </SubmenuTrigger>
                  <SubmenuTrigger>
                    <MenuItem>
                      <Plugin />
                      <Text>Skills</Text>
                    </MenuItem>
                    <Menu items={slashCommands.filter(item => item.kind === 'skill')}>
                      {item => (
                        <InsertTokenMenuItem
                          id={item.command}
                          token={{
                            type: 'token',
                            text: item.command,
                            value: {type: 'custom', anchor: '/', valueType: item.kind, data: item}
                          }}>
                          <Plugin />
                          <Text slot="label">{item.command}</Text>
                          <Text slot="description">{item.description}</Text>
                        </InsertTokenMenuItem>
                      )}
                    </Menu>
                  </SubmenuTrigger>
                  <SubmenuTrigger>
                    <MenuItem>
                      <Data />
                      <Text>Reference an object</Text>
                    </MenuItem>
                    <Menu items={objects}>
                      {item => (
                        <MenuSection>
                          <Header>
                            <Heading>{item.section}</Heading>
                          </Header>
                          <Collection items={item.items}>
                            {item => (
                              <InsertTokenMenuItem
                                id={item.title}
                                token={{
                                  type: 'token',
                                  text: item.title,
                                  value: {type: 'custom', anchor: '@', valueType: item.kind, data: item}
                                }}>
                                {item.title}
                              </InsertTokenMenuItem>
                            )}
                          </Collection>
                        </MenuSection>
                      )}
                    </Menu>
                  </SubmenuTrigger>
                </InsertMenuButton>
              </div>
              <div className={style({display: 'flex', gap: 8, alignItems: 'center'})}>
                <PromptFieldVoiceButton />
                <PromptFieldSubmitButton />
              </div>
            </PromptFieldToolbar>
          </PromptField>
        )}
      </VirtualizedStreamingChat>
    </div>
  );
}