Streamed answers with a stop button

import type { LocalMessage } from "stream-chat";

declare module "stream-chat" {
  // The field the backend sets on the assistant's replies. Optional, because every message in the
  // app shares this interface.
  interface CustomMessageData {
    ai_generated?: boolean;
  }
}

/** The assistant wrote the message: the backend marks its replies with `ai_generated`. */
export const isMessageAIGenerated = (message: LocalMessage) =>
  !!message.custom?.ai_generated;

import { useEffect, useRef } from "react";
import { AIStates } from "stream-chat";
import type { AIState } from "stream-chat";
import { useAccessibilityAnnouncer } from "stream-chat-react-native";

/** The assistant is working on an answer: thinking, checking sources or writing it. */
const isAnswering = (aiState: AIState) =>
  aiState === AIStates.Thinking ||
  aiState === AIStates.ExternalSources ||
  aiState === AIStates.Generating;

/** What to say when an answer ends in each state that isn't answering. */
const endings: Partial<Record<AIState, string>> = {
  [AIStates.Error]: "The assistant couldn't finish the answer",
  [AIStates.Idle]: "The assistant finished answering",
  [AIStates.Stop]: "Answer stopped",
};

/**
 * Announces that an answer started, and how it ended. Each change is announced once, and an
 * ending only after an answer, so a channel that opens on an earlier error stays quiet.
 */
export const useAnswerAnnouncements = (aiState: AIState) => {
  const announce = useAccessibilityAnnouncer();
  // Starts idle, so opening the channel in the middle of an answer announces it too.
  const previousState = useRef<AIState>(AIStates.Idle);

  useEffect(() => {
    if (previousState.current === aiState) return;
    const wasAnswering = isAnswering(previousState.current);
    previousState.current = aiState;

    if (isAnswering(aiState)) {
      if (!wasAnswering) announce("The assistant is answering");
    } else if (wasAnswering) {
      const ending = endings[aiState];
      if (ending) announce(ending);
    }
  }, [aiState, announce]);
};

import { useMemo } from "react";
import { ActivityIndicator, StyleSheet, Text, View } from "react-native";
import { AIStates } from "stream-chat";
import type { AIState } from "stream-chat";
import { useAIState, useChannel, useTheme } from "stream-chat-react-native";
import { useAnswerAnnouncements } from "./useAnswerAnnouncements";

/** What the status line says in each state that shows it. */
const labels: Partial<Record<AIState, string>> = {
  [AIStates.Error]: "Couldn't finish the answer",
  [AIStates.ExternalSources]: "Checking sources...",
  [AIStates.Generating]: "Writing...",
  [AIStates.Thinking]: "Thinking...",
};

/**
 * The assistant's state under the message list: a spinner and a line of text while it works, the
 * text alone after an error, and nothing otherwise. Render it inside `Channel`.
 */
export const AssistantStatus = () => {
  const channel = useChannel();
  const { aiState } = useAIState(channel);
  const {
    theme: { semantics },
  } = useTheme();
  useAnswerAnnouncements(aiState);

  const styles = useMemo(
    () =>
      StyleSheet.create({
        container: {
          alignItems: "center",
          flexDirection: "row",
          gap: 8,
          paddingHorizontal: 16,
          paddingVertical: 8,
        },
        error: { color: semantics.accentError },
        text: { color: semantics.textSecondary, fontSize: 14 },
      }),
    [semantics],
  );

  const label = labels[aiState];
  if (!label) return null;
  const isError = aiState === AIStates.Error;

  return (
    <View style={styles.container}>
      {!isError && (
        <ActivityIndicator color={semantics.accentPrimary} size="small" />
      )}
      <Text style={[styles.text, isError && styles.error]}>{label}</Text>
    </View>
  );
};

import { Button, useComponentsContext } from "stream-chat-react-native";
import type { StopMessageStreamingButtonProps } from "stream-chat-react-native";

/**
 * Fills the composer's `StopMessageStreamingButton` key, which shows in place of the send button
 * while the assistant is thinking or writing. Its `onPress` stops the answer.
 */
export const StopButton = ({ onPress }: StopMessageStreamingButtonProps) => {
  const { icons } = useComponentsContext();

  return (
    <Button
      accessibilityLabel="Stop the answer"
      iconOnly
      LeadingIcon={icons.Stop}
      onPress={onPress}
      size="sm"
      type="solid"
      variant="primary"
    />
  );
};

import type { PropsWithChildren } from "react";
import { WithComponents } from "stream-chat-react-native";
import { StopButton } from "./StopButton";

// Declared once, so WithComponents gets the same object on every render.
const overrides = { StopMessageStreamingButton: StopButton };

/** A Stop button of your own in place of the send button while the assistant answers. */
export const WithStreamedAnswers = ({ children }: PropsWithChildren) => (
  <WithComponents overrides={overrides}>{children}</WithComponents>
);

import { useMemo } from "react";
import {
  Channel,
  MessageComposer,
  MessageList,
  useChatContext,
} from "stream-chat-react-native";
import { AssistantStatus } from "./AssistantStatus";
import { WithStreamedAnswers } from "./WithStreamedAnswers";

/** The conversation with the assistant, which your backend creates for the user. */
export const AssistantChannelScreen = ({
  channelId,
}: {
  channelId: string;
}) => {
  const { client } = useChatContext();
  const channel = useMemo(
    () => client.channel("messaging", channelId),
    [channelId, client],
  );

  return (
    <WithStreamedAnswers>
      <Channel channel={channel}>
        <MessageList />
        <AssistantStatus />
        <MessageComposer />
      </Channel>
    </WithStreamedAnswers>
  );
};

import type { PropsWithChildren } from "react";
import type { StreamChat } from "stream-chat";
import { Chat, OverlayProvider } from "stream-chat-react-native";
import { isMessageAIGenerated } from "./isMessageAIGenerated";

export const App = ({
  children,
  client,
}: PropsWithChildren<{ client: StreamChat }>) => (
  <OverlayProvider>
    <Chat client={client} isMessageAIGenerated={isMessageAIGenerated}>
      {children}
    </Chat>
  </OverlayProvider>
);