add more comments

This commit is contained in:
Xuan Son Nguyen 2025-02-10 15:24:25 +01:00
parent d83a8e3731
commit 2eef3e7db4
6 changed files with 72 additions and 39 deletions

Binary file not shown.

View file

@ -13,7 +13,7 @@ interface SplitMessage {
export default function ChatMessage({
msg,
siblingLastNodeIds,
siblingLeafNodeIds,
siblingCurrIdx,
id,
onRegenerateMessage,
@ -22,7 +22,7 @@ export default function ChatMessage({
isPending,
}: {
msg: Message | PendingMessage;
siblingLastNodeIds: Message['id'][];
siblingLeafNodeIds: Message['id'][];
siblingCurrIdx: number;
id?: string;
onRegenerateMessage(msg: Message): void;
@ -45,8 +45,8 @@ export default function ChatMessage({
: null,
[msg.timings]
);
const nextSibling = siblingLastNodeIds[siblingCurrIdx + 1];
const prevSibling = siblingLastNodeIds[siblingCurrIdx - 1];
const nextSibling = siblingLeafNodeIds[siblingCurrIdx + 1];
const prevSibling = siblingLeafNodeIds[siblingCurrIdx - 1];
// for reasoning model, we split the message into content and thought
// TODO: implement this as remark/rehype plugin in the future
@ -203,7 +203,7 @@ export default function ChatMessage({
'flex-row-reverse': msg.role === 'user',
})}
>
{siblingLastNodeIds && siblingLastNodeIds.length > 1 && (
{siblingLeafNodeIds && siblingLeafNodeIds.length > 1 && (
<div className="flex gap-1 items-center opacity-60 text-sm">
<button
className={classNames({
@ -215,7 +215,7 @@ export default function ChatMessage({
<ChevronLeftIcon className="h-4 w-4" />
</button>
<span>
{siblingCurrIdx + 1} / {siblingLastNodeIds.length}
{siblingCurrIdx + 1} / {siblingLeafNodeIds.length}
</span>
<button
className={classNames({

View file

@ -6,24 +6,29 @@ import { classNames, throttle } from '../utils/misc';
import CanvasPyInterpreter from './CanvasPyInterpreter';
import StorageUtils from '../utils/storage';
/**
* A message display is a message node with additional information for rendering.
* For example, siblings of the message node are stored as their last node (aka leaf node).
*/
export interface MessageDisplay {
msg: Message | PendingMessage;
siblingLastNodeIds: Message['id'][];
siblingLeafNodeIds: Message['id'][];
siblingCurrIdx: number;
isPending?: boolean;
}
function getListMessageDisplay(
msgs: Readonly<Message[]>,
lastNodeId: Message['id']
leafNodeId: Message['id']
): MessageDisplay[] {
const currNodes = StorageUtils.filterByLastNodeId(msgs, lastNodeId, true);
const currNodes = StorageUtils.filterByLeafNodeId(msgs, leafNodeId, true);
const res: MessageDisplay[] = [];
const nodeMap = new Map<Message['id'], Message>();
for (const msg of msgs) {
nodeMap.set(msg.id, msg);
}
const findLastNode = (msgId: Message['id']): Message['id'] => {
// find leaf node from a message node
const findLeafNode = (msgId: Message['id']): Message['id'] => {
let currNode: Message | undefined = nodeMap.get(msgId);
while (currNode) {
if (currNode.children.length === 0) break;
@ -39,7 +44,7 @@ function getListMessageDisplay(
if (msg.type !== 'root') {
res.push({
msg,
siblingLastNodeIds: siblings.map(findLastNode),
siblingLeafNodeIds: siblings.map(findLeafNode),
siblingCurrIdx: siblings.indexOf(msg.id),
});
}
@ -77,11 +82,12 @@ export default function ChatScreen() {
} = useAppContext();
const [inputMsg, setInputMsg] = useState('');
// keep track of leaf node for rendering
const [currNodeId, setCurrNodeId] = useState<number>(-1);
const messages: MessageDisplay[] = useMemo(() => {
if (!viewingChat) return [];
else return getListMessageDisplay(viewingChat.messages, currNodeId);
}, [currNodeId, viewingChat?.messages]);
}, [currNodeId, viewingChat]);
const currConvId = viewingChat?.conv.id ?? null;
const pendingMsg: PendingMessage | undefined =
@ -94,7 +100,10 @@ export default function ChatScreen() {
scrollToBottom(false, 1);
}, [currConvId]);
const onChunk: CallbackGeneratedChunk = () => {
const onChunk: CallbackGeneratedChunk = (currLeafNodeId?: Message['id']) => {
if (currLeafNodeId) {
setCurrNodeId(currLeafNodeId);
}
scrollToBottom(true);
};
@ -141,12 +150,14 @@ export default function ChatScreen() {
};
const hasCanvas = !!canvasData;
// due to some timing issues of StorageUtils.appendMsg(), we need to make sure the pendingMsg is not duplicated upon rendering (i.e. appears once in the saved conversation and once in the pendingMsg)
const pendingMsgDisplay: MessageDisplay[] =
pendingMsg && messages.at(-1)?.msg.id !== pendingMsg.id
? [
{
msg: pendingMsg,
siblingLastNodeIds: [],
siblingLeafNodeIds: [],
siblingCurrIdx: 0,
isPending: true,
},
@ -178,7 +189,7 @@ export default function ChatScreen() {
<ChatMessage
key={msg.msg.id}
msg={msg.msg}
siblingLastNodeIds={msg.siblingLastNodeIds}
siblingLeafNodeIds={msg.siblingLeafNodeIds}
siblingCurrIdx={msg.siblingCurrIdx}
onRegenerateMessage={handleRegenerateMessage}
onEditMessage={handleEditMessage}

View file

@ -23,7 +23,7 @@ interface AppContextValue {
isGenerating: (convId: string) => boolean;
sendMessage: (
convId: string | null,
lastNodeId: Message['id'] | null,
leafNodeId: Message['id'] | null,
content: string,
onChunk: CallbackGeneratedChunk
) => Promise<boolean>;
@ -47,7 +47,7 @@ interface AppContextValue {
}
// this callback is used for scrolling to the bottom of the chat and switching to the last node
export type CallbackGeneratedChunk = () => void;
export type CallbackGeneratedChunk = (currLeafNodeId?: Message['id']) => void;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const AppContext = createContext<AppContextValue>({} as any);
@ -130,7 +130,7 @@ export const AppContextProvider = ({
const generateMessage = async (
convId: string,
lastNodeId: Message['id'],
leafNodeId: Message['id'],
onChunk: CallbackGeneratedChunk
) => {
if (isGenerating(convId)) return;
@ -141,9 +141,9 @@ export const AppContextProvider = ({
throw new Error('Current conversation is not found');
}
const currMessages = StorageUtils.filterByLastNodeId(
const currMessages = StorageUtils.filterByLeafNodeId(
await StorageUtils.getMessages(convId),
lastNodeId,
leafNodeId,
false
);
const abortController = new AbortController();
@ -161,7 +161,7 @@ export const AppContextProvider = ({
timestamp: pendingId,
role: 'assistant',
content: null,
parent: lastNodeId,
parent: leafNodeId,
children: [],
};
setPending(convId, pendingMsg);
@ -264,26 +264,26 @@ export const AppContextProvider = ({
}
if (pendingMsg.content !== null) {
await StorageUtils.appendMsg(pendingMsg as Message, lastNodeId);
await StorageUtils.appendMsg(pendingMsg as Message, leafNodeId);
}
setPending(convId, null);
onChunk(); // trigger scroll to bottom and switch to the last node
onChunk(pendingId); // trigger scroll to bottom and switch to the last node
};
const sendMessage = async (
convId: string | null,
lastNodeId: Message['id'] | null,
leafNodeId: Message['id'] | null,
content: string,
onChunk: CallbackGeneratedChunk
): Promise<boolean> => {
if (isGenerating(convId ?? '') || content.trim().length === 0) return false;
if (convId === null || convId.length === 0 || lastNodeId === null) {
if (convId === null || convId.length === 0 || leafNodeId === null) {
const conv = await StorageUtils.createConversation(
content.substring(0, 256)
);
convId = conv.id;
lastNodeId = conv.currNode;
leafNodeId = conv.currNode;
// if user is creating a new conversation, redirect to the new conversation
navigate(`/chat/${convId}`);
}
@ -298,12 +298,12 @@ export const AppContextProvider = ({
convId,
role: 'user',
content,
parent: lastNodeId,
parent: leafNodeId,
children: [],
},
lastNodeId
leafNodeId
);
onChunk();
onChunk(currMsgId);
try {
await generateMessage(convId, currMsgId, onChunk);
@ -346,8 +346,8 @@ export const AppContextProvider = ({
);
parentNodeId = currMsgId;
}
onChunk(parentNodeId);
onChunk();
await generateMessage(convId, parentNodeId, onChunk);
};

View file

@ -48,17 +48,17 @@ const StorageUtils = {
return (await db.conversations.where('id').equals(convId).first()) ?? null;
},
/**
* get messages by convId and timeline
* get all message nodes in a conversation
*/
async getMessages(convId: string): Promise<Message[]> {
return await db.messages.where({ convId }).toArray();
},
/**
* use in conjunction with getMessages to filter messages by lastNodeId
* use in conjunction with getMessages to filter messages by leafNodeId
*/
filterByLastNodeId(
filterByLeafNodeId(
msgs: Readonly<Message[]>,
lastNodeId: Message['id'],
leafNodeId: Message['id'],
includeRoot: boolean
): Readonly<Message[]> {
const res: Message[] = [];
@ -66,7 +66,7 @@ const StorageUtils = {
for (const msg of msgs) {
nodeMap.set(msg.id, msg);
}
let startNode: Message | undefined = nodeMap.get(lastNodeId);
let startNode: Message | undefined = nodeMap.get(leafNodeId);
if (!startNode) {
// if not found, we return the path with the latest timestamp
let latestTime = -1;
@ -77,7 +77,7 @@ const StorageUtils = {
}
}
}
// traverse the path from lastNodeId to root
// traverse the path from leafNodeId to root
// startNode can never be undefined here
let currNode: Message | undefined = startNode;
while (currNode) {
@ -89,7 +89,7 @@ const StorageUtils = {
return res;
},
/**
* create a new conversation with a default timeline number 0
* create a new conversation with a default root node
*/
async createConversation(name: string): Promise<Conversation> {
const now = Date.now();

View file

@ -7,11 +7,33 @@ export interface TimingReport {
/**
* What is conversation "branching"? It is a feature that allows the user to edit an old message in the history, while still keeping the conversation flow.
* Inspired by ChatGPT UI where you edit a message, a new branch of the conversation is created, and the old message is still visible.
* Inspired by ChatGPT / Claude / Hugging Chat where you edit a message, a new branch of the conversation is created, and the old message is still visible.
*
* We use the same node based structure as ChatGPT, where each message has a parent and children. A "root" message is the first message in a conversation, which will not be displayed in the UI.
* We use the same node-based structure like other chat UIs, where each message has a parent and children. A "root" message is the first message in a conversation, which will not be displayed in the UI.
*
* root
* message 1
* message 2
* message 3
* message 4
* message 5
*
* In the above example, assuming that user wants to edit message 2, a new branch will be created:
*
* message 2
* message 3
* message 6
*
* Message 2 and 6 are siblings, and message 6 is the new branch.
*
* We only need to know the last node (aka leaf) to get the current branch. In the above example, message 5 is the leaf of branch containing message 4 and 5.
*
* For the implementation:
* - StorageUtils.getMessages() returns list of all nodes
* - StorageUtils.filterByLeafNodeId() filters the list of nodes from a given leaf node
*/
// Note: the term "message" and "node" are used interchangeably in this context
export interface Message {
id: number;
convId: string;