feat: change messagesAtom to a WritableAtom

jotaijs · Nov 20, 2023 · 5c1cf03 · 5c1cf03
1 parent fd72b2d
commit 5c1cf03
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -6,7 +6,7 @@ with [Vercel AI SDK](https://sdk.vercel.ai/docs).
 ## install
 
 ```
-yarn add jotai-ai
+yarn add ai jotai-ai
 ```
 
 ## chatAtoms
@@ -77,10 +77,28 @@ They actually have the different behaviors in the different frameworks.
 `chatAtoms` provider a more flexible way to create the chatbot, which is based on `jotai` atoms, so you can use it in
 framework-agnostic way.
 
+For example, you can extend the messagesAtom to add more features like `clearMessagesAtom`:
+
+```js
+const { messagesAtom } = chatAtoms()
+
+const clearMessagesAtom = atom(
+  null,
+  async (get, set) => set(messagesAtom, [])
+)
+
+const Actions = () => {
+  const clear = useSetAtom(clearMessagesAtom)
+  return (
+    <button onClick={clear}>Clear Messages</button>
+  )
+}
+```
+
 Also, `chatAtoms` is created out of the Component lifecycle,
 so you can share the state between different components easily.
 
-#### Load messages on demand
+#### Load messages on demand with React Suspense
 
 `chatAtoms` also allows you to pass async fetch function to `initialMessage` option, which is not supported by `useChat`.
 
@@ -101,11 +119,6 @@ const {
 With the combination with `jotai-effect`, you can create a chatbot with local storage support.
 
 ```js
-import { Suspense } from 'react'
-import { useAtomValue } from 'jotai'
-import { chatAtoms } from 'jotai-ai'
-import { atomEffect } from 'jotai-effect'
-
 const {
   messagesAtom
 } = chatAtoms({

diff --git a/examples/llamaindex-straming/app/components/chat-section.tsx b/examples/llamaindex-straming/app/components/chat-section.tsx
@@ -6,6 +6,7 @@ import { atomEffect } from 'jotai-effect'
 import { ChatInput, ChatMessages } from './ui/chat'
 import { useAtom, useAtomValue, useSetAtom } from 'jotai/react'
 import { Suspense } from 'react'
+import { atom } from 'jotai/vanilla'
 
 const {
   messagesAtom,
@@ -22,6 +23,11 @@ const {
   }
 })
 
+const clearMessagesAtom = atom(
+  null,
+  async (get, set) => set(messagesAtom, [])
+)
+
 const saveMessagesEffectAtom = atomEffect((get, set) => {
   const messages = get(messagesAtom)
   const idbPromise = import('idb-keyval')
@@ -40,12 +46,14 @@ const saveMessagesEffectAtom = atomEffect((get, set) => {
 const Messages = () => {
   const messages = useAtomValue(messagesAtom)
   const isLoading = useAtomValue(isLoadingAtom)
+  const clear = useSetAtom(clearMessagesAtom)
   const reload = useSetAtom(reloadAtom)
   const stop = useSetAtom(stopAtom)
   return (
     <ChatMessages
       messages={messages}
       isLoading={isLoading}
+      clear={clear}
       reload={reload}
       stop={stop}
     />

diff --git a/examples/llamaindex-straming/app/components/ui/chat/chat-actions.tsx b/examples/llamaindex-straming/app/components/ui/chat/chat-actions.tsx
@@ -1,16 +1,23 @@
-import { PauseCircle, RefreshCw } from "lucide-react";
+import { PauseCircle, RefreshCw, Eraser } from "lucide-react";
 
 import { Button } from "../button";
 import { ChatHandler } from "./chat.interface";
 
 export default function ChatActions(
-  props: Pick<ChatHandler, "stop" | "reload"> & {
+  props: Pick<ChatHandler, "stop" | "reload" | "clear"> & {
+    showClear?: boolean;
     showReload?: boolean;
     showStop?: boolean;
   },
 ) {
   return (
     <div className="space-x-4">
+      {props.showClear && (
+        <Button variant="outline" size="sm" onClick={props.clear}>
+          <Eraser className="mr-2 h-4 w-4" />
+          Clear
+        </Button>
+      )}
       {props.showStop && (
         <Button variant="outline" size="sm" onClick={props.stop}>
           <PauseCircle className="mr-2 h-4 w-4" />

diff --git a/examples/llamaindex-straming/app/components/ui/chat/chat-messages.tsx b/examples/llamaindex-straming/app/components/ui/chat/chat-messages.tsx
@@ -5,7 +5,7 @@ import ChatMessage from "./chat-message";
 import { ChatHandler } from "./chat.interface";
 
 export default function ChatMessages(
-  props: Pick<ChatHandler, "messages" | "isLoading" | "reload" | "stop">,
+  props: Pick<ChatHandler, "messages" | "isLoading" | "reload" | "stop" | "clear">,
 ) {
   const scrollableChatContainerRef = useRef<HTMLDivElement>(null);
   const messageLength = props.messages.length;
@@ -20,6 +20,7 @@ export default function ChatMessages(
 
   const isLastMessageFromAssistant =
     messageLength > 0 && lastMessage?.role !== "user";
+  const showClear = props.clear && messageLength > 0;
   const showReload =
     props.reload && !props.isLoading && isLastMessageFromAssistant;
   const showStop = props.stop && props.isLoading;
@@ -40,8 +41,10 @@ export default function ChatMessages(
       </div>
       <div className="flex justify-end py-4">
         <ChatActions
+          clear={props.clear}
           reload={props.reload}
           stop={props.stop}
+          showClear={showClear}
           showReload={showReload}
           showStop={showStop}
         />

diff --git a/examples/llamaindex-straming/app/components/ui/chat/chat.interface.ts b/examples/llamaindex-straming/app/components/ui/chat/chat.interface.ts
@@ -10,6 +10,7 @@ export interface ChatHandler {
   isLoading: boolean;
   handleSubmit: (e: React.FormEvent<HTMLFormElement>) => void;
   handleInputChange: (e: React.ChangeEvent<HTMLInputElement>) => void;
+  clear?: () => void;
   reload?: () => void;
   stop?: () => void;
 }
diff --git a/src/index.ts b/src/index.ts
@@ -1,4 +1,4 @@
-'use client';
+'use client'
 import { atom, type Getter, type Setter } from 'jotai/vanilla'
 import type {
   ChatRequest,
@@ -14,6 +14,15 @@ import { parseComplexResponse } from './parse-complex-response'
 
 const COMPLEX_HEADER = 'X-Experimental-Stream-Data'
 
+const isPromiseLike = (value: unknown): value is PromiseLike<unknown> => {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    'then' in value &&
+    typeof (value as any).then === 'function'
+  )
+}
+
 export function chatAtoms (
   chatOptions: Omit<
     UseChatOptions,
@@ -45,7 +54,7 @@ export function chatAtoms (
   )
   const messagesAtom = atom<
     Message[] | Promise<Message[]>,
-    [messages: Message[]],
+    [messages: ((Message[]) | Promise<Message[]>)],
     void
   >(
     get => {
@@ -174,7 +183,7 @@ export function chatAtoms (
 
         // The request has been aborted, stop reading the stream.
         if (abortController.signal.aborted) {
-          reader.cancel()
+          await reader.cancel()
           break
         }
       }
@@ -350,7 +359,20 @@ export function chatAtoms (
 
   // user side atoms
   return {
-    messagesAtom: atom(get => get(messagesAtom)),
+    messagesAtom: atom(
+      get => get(messagesAtom),
+      async (
+        get,
+        set,
+        messages: Message[]
+      ): Promise<void> => {
+        const prevMessages = get(messagesAtom)
+        if (isPromiseLike(prevMessages)) {
+          set(messagesAtom, prevMessages.then(() => messages))
+        } else {
+          set(messagesAtom, messages)
+        }
+      }),
     isLoadingAtom: atom(get => get(isLoadingAtom)),
     inputAtom: atom(
       get => get(inputBaseAtom),