feature: Refactor jotai-ai to fit new architecture since AI SDK v5 #20
Conversation
himself65
left a comment
himself65
left a comment
There was a problem hiding this comment.
I feel like this is good. Thank u for doing the migration. It's okay to break the things for v5
|
lmk if you have any updates |
Ok. I see. Trying to focus on this asap 🥹 |
|
is syncing upon the like this: export function useJotaiChat({
chat: Chat,
...opts
}) {
const [messages, setMessages] = useAtom(messageAtom);
const [status, setStatus] = useAtom(statusAtom);
const [error, setError] = useAtom(errorAtom);
useEffect(() => {
const unregisterMessagesCallback = chat[
"~registerMessagesCallback"
](() => {
setMessages(chat.messages);
}, 50);
const unregisterStatusCallback = chat["~registerStatusCallback"](
() => {
setStatus(chat.status);
},
);
const unregisterErrorCallback = chat["~registerErrorCallback"](
() => {
setError(chat.error);
},
);
return () => {
unregisterMessagesCallback();
unregisterStatusCallback();
unregisterErrorCallback();
};
}, [chat, setMessages, setStatus, setError]);
const ret = useChat({
chat,
...opts
});
return {
...ret,
messages,
status,
error,
};
} |
|
Thank @songkeys for his innovated solutions. I tried it, and it actually works in some test cases, I'm also grad to learn something new from it! Unfortunately, that's not enough. Then say hello to all subscribers, I'm here to sharing my progress recently. My first try was structurally inspired (emmmm, almost copied) from FranciscoMoretti's Gist - Zustand + AI SDK useChat Integration. However, it just "looks great". When I added unit tests, my first version failed everywhere (for example, infinite re-render loops). Then I started to debug and dive into the internal details of The new I was thinking, the key of the solution is how to cooperate I repost the most enlightening illumination here TL;DR:
Hence, my original schedules would be:
PS: I'm now trying to play some magical/dark tricks via jotai Now, let me clarify why @songkeys' approach is not enough:
But I indeed get inspired from @songkeys' suggestions. Especially, I have new ideas and new directions when I'm writing and organizing this reply (aha, that's how discussions work 🤣). My eventual desired goals of
So the API will be like: const initialMessagesAtom = atom(...);
const initialChatId = atom(...);
// Just to show the most jotai primitive function call
// const { messagesAtom, statusAtom, errorAtom } = atomAIChat(
// get => {},
// (get, set, ...) => {},
// );
const { messagesAtom, statusAtom, errorAtom } = atomWithAIChat({
initialMessagesAtom
});
// So what the `jotai-ai` really do is to help end developers to write getter and setter.
// Keep the dark detail inside of `atomWithAIChat`
// and let end users to compose or derive atoms as they hope!.
const ChatComponent = () => {
const [messages, setMessages] = useAtom(messagesAtom)
// Or actually we could customize a hook to expose more handlers
const {messages, sendMessages, resendMessages, abortSendingMessages} = useChatMessagesAtom(messagesAtom)
return (<>
...
...
<>)
}IMHO, that's the "Jotai style", that's also my motivation of PR #15 which I look like to forget yet. What I have mistaken is I hope to reuse So next step, I plan to try some new API signatures to improve overall usabilities. Any further discussion and suggestion is appreciated and welcome! See what amazing we have gained now! Thanks for all! |
|
Thank you for sharing your thoughts! I learned a lot from your insights. Your new design truly embodies the jotai approach. Here's a bit more about my current usage in my project, as I mentioned earlier: Most AI chat applications today support multiple chats, and mine is no exception. I define my chat atoms like this: export const chatMessagesFamily = atomFamily((chatId: string) => atom<Chat["messages"]>([]));
export const chatStatusFamily = atomFamily((chatId: string) => atom<Chat["status"]>("ready"));
export const chatErrorFamily = atomFamily((chatId: string) => atom<Chat["error"]>(undefined));Next, I create a global const globalChatInstance: Record<string, Chat> = {};
const getChatInstance = (chatId: string) => {
if (globalChatInstance[chatId]) {
return globalChatInstance[chatId];
}
const chatInstance = new Chat({
id: chatId,
onError: (error) => {
// We can manipulate the `messages` in the current chatInstance directly like this
chatInstance.messages = chatInstance.messages.slice(0, -1);
// it will sync with jotai atoms
},
onData: (part) => {
// We can also set the jotai store
// For example, here we update the current chat's title (which is stored in another atom)
if (part.type === "data-title") {
const nextTitle = part.data.text;
store.set(chatConversationsAtom, (prev) => {
return prev.map((c) =>
c.id === chatId ? { ...c, title: nextTitle } : c,
);
});
}
},
});
globalChatInstance[chatId] = chatInstance;
return chatInstance;
};Then, the export function useJotaiChat({
chatId: string
}) {
const chatInstance = getChatInstance(chatId);
const [messages, setMessages] = useAtom(chatMessagesFamily(chatId));
const [status, setStatus] = useAtom(chatStatusFamily(chatId));
const [error, setError] = useAtom(chatErrorFamily(chatId));
useEffect(() => {
const unregisterMessagesCallback = chat[
"~registerMessagesCallback"
](() => {
setMessages(chat.messages);
}, 50);
const unregisterStatusCallback = chat["~registerStatusCallback"](
() => {
setStatus(chat.status);
},
);
const unregisterErrorCallback = chat["~registerErrorCallback"](
() => {
setError(chat.error);
},
);
return () => {
unregisterMessagesCallback();
unregisterStatusCallback();
unregisterErrorCallback();
};
}, [chat, setMessages, setStatus, setError]);
const ret = useChat({
chat,
...opts
});
return {
...ret,
messages,
status,
error,
};
}You mentioned:
Yes. It will do.
Due to the global singletons, I actually use it directly in my child components. It might not be the most efficient way, but it feels natural to me so far. The only downside is that Of course, I can also use
Yes. You're right. I need to only use I have to admit that my implementation isn't quite satisfying. I haven't thought deeply about it; I've just been trying to get my project completed... And when I realized I should have searched for something like Back to your new design, I believe we're on the right track. const chatIdAtom = atom("random-id");
const chatAtom = atomWithChat((get) => ({
id: get(chatIdAtom)
}));
const [{ messages, sendMessage, stop }] = useAtom(chatAtom); |
|
@songkeys Thanks for sharing your use cases and feedbacks! That's great and also what we want.
Sure, let's try together and do experiments to find the best approaches. |
|
@songkeys I have updated API design. Please check, does the new API fit your use cases :) @himself65 It seems I have done the major design, you could do basic review now? I will keep polishing the rest parts, e.g.: tune API options, fix corner cases and add new docs. (It's very embarrassed that all tests passed in my local env, but become flaky in CI. Sometime successes, sometime fails. I haven't figured out why) Please see the first update section in the main description. |
3 participants
I have tried AI SDK v5 + Jotai in my personal project to leverage new architecture. It's indeed much easier to use and maintain.
Current PR is still draft. @himself65 you could have a quick look for new implementation under
v5folder.And I hope some suggestions like do I need to remove all legacy
make-chat-atoms? Or we could open a new branch to release newer package?Extra actions
First update
After discussion with @songkeys, I have updated the design of API. Adding one atom constructor
atomWithChatand one hook functionuseChatAtomValue. Basically, try to reuseAbstractChatand rewriteuseChatto match the philosophy of Jotai itself.First, declare the
chatAtomthen, we could use the
chatAtomin the component like other atoms with customized hookuseChatAtomValueYes, the API is very similar to the
jotai-tanstack-query, the limitation parts we will discuss later.For the first input argument of the
atomWithChat, it's just a reader function, so we could compose other atoms and state easily,Compatible with SSR and whatever other approaches to initial ID or messages by Jotai style atoms (like
atomFamily).That's all! No more dark magic, just Jotai.
Current results of UI tests (adapt
useChattest cases from upstream) are added as follows:Note
The stderr output "Not implemented: HTMLFormElement's requestSubmit() method" is not our responsibility.
I have passed almost all UI tests, except one corner case, which may fix later.
What extra misc I added in this PR:
full-testforpushandpull_requestbuildscript tobuildandbuild:examplesto improve scope of building.Known issues and limitations
useChatAtomValueis forcedly required now.Unlike
jotai-tanstack-query, we cannot use nativejotaihooks likeuseAtom/useAtomValue(useSetAtomis fine).Unfortunately,⚠️ The above two hooks have different behaviors now. The good news is the caused reason is clear and fixable. In general, we have no
atomWithExternalSyncoratomWithObservable, so we cannot subscribe change events from external sources insideatomWithChat. Hence, we have to implement the sync effect insideuseChatAtomValueoruseChatAtom, see the following example:That's why they have different behaviors now. If we can move the sync effect logic into
atomWithChatviaatomWithExternalStorage, then we unify them in the future.I also found some discussions to implement
atomWithExternalStoragein 2022, but have no more progress since then.Current solution: We can try to add explicit warnings in docs to tell users to avoid mistakes.
setMessages(handler inuseChatreturns) yet.As ai-sdk doc say:
But actually, IMHO, it's quite unusable or unsafe operation. There is no more mechanism to deal with what if optimistic updates failed even use the official API. The user still need to implement all other details. So I think we should left this part to the user self, if they know what they're doing,
setMessagesis not difficult to implement withjotai-ai