Skip to content

Commit 46b1a76

Browse files
MarcMcIntoshMarcMcIntosh
committed
wip(chat tree): render a basic tree what can be navigated through.
1 parent d36db1f commit 46b1a76

File tree

9 files changed

+218
-93
lines changed

9 files changed

+218
-93
lines changed

‎refact-agent/gui/src/components/ChatContent/UserInput.tsx‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,10 @@ import styles from "./ChatContent.module.css";
1818

1919
export type UserInputProps = {
2020
children: UserMessage["content"];
21-
messageIndex: number;
21+
// TODO: remove when using nodes
22+
messageIndex?: number;
2223
// maybe add images argument ?
23-
onRetry: (index: number, question: UserMessage["content"]) => void;
24+
onRetry?: (index: number, question: UserMessage["content"]) => void;
2425
// disableRetry?: boolean;
2526
};
2627

@@ -29,15 +30,15 @@ export const UserInput: React.FC<UserInputProps> = ({
2930
children,
3031
onRetry,
3132
}) => {
32-
const messages = useAppSelector(selectMessages);
33+
// const messages = useAppSelector(selectMessages);
3334

3435
const [showTextArea, setShowTextArea] = useState(false);
3536
const [isEditButtonVisible, setIsEditButtonVisible] = useState(false);
3637
// const ref = React.useRef<HTMLButtonElement>(null);
3738

3839
const handleSubmit = useCallback(
3940
(value: UserMessage["content"]) => {
40-
onRetry(messageIndex, value);
41+
onRetry && messageIndex && onRetry(messageIndex, value);
4142
setShowTextArea(false);
4243
},
4344
[messageIndex, onRetry],
@@ -58,11 +59,13 @@ export const UserInput: React.FC<UserInputProps> = ({
5859
const isString = typeof children === "string";
5960
const linesLength = isString ? children.split("\n").length : Infinity;
6061

61-
const checkpointsFromMessage = useMemo(() => {
62-
const maybeUserMessage = messages[messageIndex];
63-
if (!isUserMessage(maybeUserMessage)) return null;
64-
return maybeUserMessage.checkpoints;
65-
}, [messageIndex, messages]);
62+
// TODO: add this back in
63+
// const checkpointsFromMessage = useMemo(() => {
64+
// if (!messageIndex) return null;
65+
// const maybeUserMessage = messages[messageIndex];
66+
// if (!isUserMessage(maybeUserMessage)) return null;
67+
// return maybeUserMessage.checkpoints;
68+
// }, [messageIndex, messages]);
6669

6770
return (
6871
<Container position="relative" pt="1">
@@ -105,12 +108,12 @@ export const UserInput: React.FC<UserInputProps> = ({
105108
transition: "opacity 0.15s, visibility 0.15s",
106109
}}
107110
>
108-
{checkpointsFromMessage && checkpointsFromMessage.length > 0 && (
111+
{/* {checkpointsFromMessage && checkpointsFromMessage.length > 0 && (
109112
<CheckpointButton
110113
checkpoints={checkpointsFromMessage}
111114
messageIndex={messageIndex}
112115
/>
113-
)}
116+
)} */}
114117
<IconButton
115118
title="Edit message"
116119
variant="soft"

‎refact-agent/gui/src/components/ChatHistory/ChatHistory.tsx‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import React, { memo, useCallback, useEffect, useMemo } from "react";
1+
import React, { useCallback, useEffect } from "react";
22
import { Flex, Box } from "@radix-ui/themes";
33
import { ScrollArea } from "../ScrollArea";
44
import { HistoryItem } from "./HistoryItem";
@@ -33,10 +33,15 @@ function useGetHistory() {
3333
});
3434
const isLoading = useAppSelector(chatDbSelectors.getLoading);
3535

36+
// move this to a dedicated hook
3637
useEffect(() => {
3738
const thunk = dispatch(subscribeToThreadsThunk());
3839
return () => {
39-
thunk.abort("unmounted");
40+
try {
41+
thunk.abort("unmounted");
42+
} catch {
43+
// noop
44+
}
4045
};
4146
}, [dispatch]);
4247

‎refact-agent/gui/src/components/MessageNode/MessageNode.stories.tsx‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import React from "react";
2+
import type { Meta, StoryObj } from "@storybook/react";
3+
import { MessageNode } from "./MessageNode";
4+
import { CMESSAGES_WITH_NESTED_BRANCHES_STUB } from "../../__fixtures__";
5+
import { makeMessageTree } from "../../features/ChatDB/makeMessageTree";
6+
import { Provider } from "react-redux";
7+
import { Theme } from "../Theme";
8+
import { AbortControllerProvider } from "../../contexts/AbortControllers";
9+
import { setUpStore } from "../../app/store";
10+
import { CMessageNode } from "../../features/ChatDB/chatDbMessagesSlice";
11+
12+
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
13+
const messageTree = makeMessageTree(CMESSAGES_WITH_NESTED_BRANCHES_STUB)!;
14+
15+
const Template: React.FC<{ node: CMessageNode }> = ({ node }) => {
16+
const store = setUpStore();
17+
18+
return (
19+
<Provider store={store}>
20+
<Theme>
21+
<AbortControllerProvider>
22+
<MessageNode>{node}</MessageNode>
23+
</AbortControllerProvider>
24+
</Theme>
25+
</Provider>
26+
);
27+
};
28+
const meta: Meta<typeof Template> = {
29+
title: "components/MessageNode",
30+
component: Template,
31+
};
32+
33+
export default meta;
34+
35+
export const Primary: StoryObj<typeof Template> = {
36+
args: { node: messageTree },
37+
};

‎refact-agent/gui/src/components/MessageNode/MessageNode.tsx‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import React from "react";
2+
import {
3+
CMessageNode,
4+
isUserCMessageNode,
5+
UserCMessageNode,
6+
} from "../../features/ChatDB/chatDbMessagesSlice";
7+
import { UserInput } from "../ChatContent/UserInput";
8+
import { AssistantInput } from "../ChatContent/AssistantInput";
9+
import { ChatMessage } from "../../services/refact";
10+
import { IconButton } from "@radix-ui/themes";
11+
import { ArrowLeftIcon, ArrowRightIcon } from "@radix-ui/react-icons";
12+
13+
const ElementForNodeMessage: React.FC<{ message: ChatMessage }> = ({
14+
message,
15+
}) => {
16+
if (message.role === "user") {
17+
return <UserInput>{message.content}</UserInput>;
18+
}
19+
20+
if (message.role === "assistant") {
21+
return (
22+
<AssistantInput
23+
message={message.content}
24+
toolCalls={message.tool_calls}
25+
/>
26+
);
27+
}
28+
29+
return false;
30+
};
31+
32+
export type MessageNodeProps = { children?: CMessageNode };
33+
export const MessageNode: React.FC<MessageNodeProps> = ({ children }) => {
34+
if (!children) return null;
35+
return (
36+
<>
37+
<ElementForNodeMessage message={children.message.cmessage_json} />
38+
<MessageNodeChildren>{children.children}</MessageNodeChildren>
39+
</>
40+
);
41+
};
42+
43+
const MessageNodeChildren: React.FC<{ children: CMessageNode[] }> = ({
44+
children,
45+
}) => {
46+
const userMessages: UserCMessageNode[] = children.filter(isUserCMessageNode);
47+
48+
if (userMessages.length === 0) {
49+
return children.map((node, index) => {
50+
const key = `${node.message.cmessage_belongs_to_cthread_id}_${node.message.cmessage_num}_${node.message.cmessage_alt}_${index}`;
51+
return <MessageNode key={key}>{node}</MessageNode>;
52+
});
53+
} else {
54+
return <UserMessageNode>{userMessages}</UserMessageNode>;
55+
}
56+
};
57+
58+
const UserMessageNode: React.FC<{ children: UserCMessageNode[] }> = ({
59+
children,
60+
}) => {
61+
// info about the node may need to be shared with the user input
62+
const [selectedNodeIndex, setSelectedNodeIndex] = React.useState<number>(0);
63+
64+
const selectedNode = children[selectedNodeIndex];
65+
return (
66+
<>
67+
<IconButton
68+
variant="outline"
69+
size="1"
70+
disabled={selectedNodeIndex === 0}
71+
onClick={() =>
72+
setSelectedNodeIndex((prev) => {
73+
if (prev === 0) return prev;
74+
return prev - 1;
75+
})
76+
}
77+
>
78+
<ArrowLeftIcon />
79+
</IconButton>
80+
<IconButton
81+
variant="outline"
82+
size="1"
83+
disabled={selectedNodeIndex === children.length - 1}
84+
onClick={() => {
85+
setSelectedNodeIndex((prev) => {
86+
if (prev === children.length - 1) return prev;
87+
return prev + 1;
88+
});
89+
}}
90+
>
91+
<ArrowRightIcon />
92+
</IconButton>
93+
<UserInput>{selectedNode.message.cmessage_json.content}</UserInput>
94+
<MessageNodeChildren>{selectedNode.children}</MessageNodeChildren>
95+
</>
96+
);
97+
};

‎refact-agent/gui/src/components/MessageNode/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export * from "./MessageNode";

‎refact-agent/gui/src/features/ChatDB/chatDbMessagesSlice.ts‎

Lines changed: 19 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,32 @@ import {
55
CThreadDefault,
66
CMessage,
77
ChatMessage,
8+
UserCMessage,
9+
isUserCMessage,
810
} from "../../services/refact";
911
import { v4 as uuid } from "uuid";
1012
import { parseOrElse } from "../../utils";
13+
import { makeMessageTree } from "./makeMessageTree";
1114

12-
export type CMessageNode = {
15+
export interface CMessageNode {
1316
message: CMessage;
1417
children: CMessageNode[];
15-
};
18+
}
1619

1720
export type CMessageRoot = CMessageNode[];
1821

22+
export interface UserCMessageNode extends CMessageNode {
23+
message: UserCMessage;
24+
}
25+
26+
export function isUserCMessageNode(
27+
node: CMessageNode,
28+
): node is UserCMessageNode {
29+
return isUserCMessage(node.message);
30+
}
31+
1932
type InitialState = {
2033
thread: CThread | CThreadDefault;
21-
messageTree: CMessageRoot;
2234
messageList: CMessage[];
2335
loading: boolean;
2436
error: null | string;
@@ -37,30 +49,11 @@ const createChatThread = (): CThreadDefault => {
3749

3850
const initialState: InitialState = {
3951
thread: createChatThread(),
40-
messageTree: [[]],
4152
messageList: [],
4253
loading: false,
4354
error: null,
4455
};
4556

46-
const findNodeByAltAndNum = (
47-
nodes: CMessageNode[],
48-
alt: number,
49-
num: number,
50-
): CMessageNode | null => {
51-
for (const node of nodes) {
52-
if (
53-
node.message.cmessage_alt === alt &&
54-
node.message.cmessage_num === num
55-
) {
56-
return node;
57-
}
58-
const found = findNodeByAltAndNum(node.children, alt, num);
59-
if (found) return found;
60-
}
61-
return null;
62-
};
63-
6457
function parseCMessageFromChatDBToCMessage(
6558
message: CMessageFromChatDB,
6659
): CMessage | null {
@@ -79,7 +72,6 @@ export const chatDbMessageSlice = createSlice({
7972
reducers: {
8073
setThread: (state, action: PayloadAction<CThread>) => {
8174
state.thread = action.payload;
82-
state.messageTree = [];
8375
},
8476
updateMessage: (
8577
state,
@@ -88,64 +80,14 @@ export const chatDbMessageSlice = createSlice({
8880
if (action.payload.threadId !== state.thread.cthread_id) return state;
8981
const message = parseCMessageFromChatDBToCMessage(action.payload.message);
9082
if (!message) return;
91-
9283
// Update message list
9384
state.messageList[message.cmessage_num] = message;
94-
95-
if (message.cmessage_num === 0) {
96-
state.messageTree[message.cmessage_num] = {
97-
message,
98-
children: state.messageTree[message.cmessage_num]?.children ?? [],
99-
};
100-
return;
101-
}
102-
103-
// find it's place
104-
// function traverse(node: CMessageNode) {}
105-
// updateMessage: (
106-
// state,
107-
// action: PayloadAction<{ threadId: string; message: CMessageFromChatDB }>,
108-
// ) => {
109-
// if (action.payload.threadId !== state.thread.cthread_id) return state;
110-
// const message = parseCMessageFromChatDBToCMessage(action.payload.message);
111-
// if (!message) return;
112-
113-
// state.messageList[message.cmessage_num] = message;
114-
// if (message.cmessage_num === 0) {
115-
// state.messageTree[message.cmessage_num] = {
116-
// message,
117-
// children: state.messageTree[message.cmessage_num]?.children ?? [],
118-
// };
119-
120-
// return;
121-
// }
122-
123-
// // find the parent node,
124-
// const parentMessage = state.messageList.find(
125-
// (m) =>
126-
// m.cmessage_num === message.cmessage_num - 1 &&
127-
// m.cmessage_alt === message.cmessage_prev_alt,
128-
// );
129-
// console.log("parentMessage", JSON.stringify(parentMessage)
130-
131-
// // const parentNode = findNodeByAltAndNum(
132-
// // state.messageTree.flat(),
133-
// // prevAlt,
134-
// // prevNum
135-
// // );
136-
137-
// // state.messageTree[message.cmessage_num] =
138-
// // state.messageTree[message.cmessage_num] ?? [];
139-
// // state.messageTree[message.cmessage_num][message.cmessage_alt] = {
140-
// // children:
141-
// // state.messageTree[message.cmessage_num][message.cmessage_alt]
142-
// // ?.children ?? [],
143-
// // message,
144-
// // };
145-
// // const node = row[message.cmessage_alt] ?? { message, children: [] };
146-
// },
14785
},
14886
},
87+
88+
selectors: {
89+
selectMessageTree: (state) => makeMessageTree(state.messageList),
90+
},
14991
});
15092

15193
export const chatDbMessageSliceActions = chatDbMessageSlice.actions;

‎refact-agent/gui/src/features/ChatDB/makeMessageTree.test.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { expect, describe, test } from "vitest";
2-
import { CMessage } from "../../services/refact";
32
import { CMessageNode } from "./chatDbMessagesSlice";
4-
import { makeMessageTree } from "./makeMessageTrie";
3+
import { makeMessageTree } from "./makeMessageTree";
54
import {
65
CMESSAGES_STUB,
76
CMESSAGES_WITH_NESTED_BRANCHES_STUB,

‎refact-agent/gui/src/features/ChatDB/makeMessageTrie.ts‎ renamed to ‎refact-agent/gui/src/features/ChatDB/makeMessageTree.ts‎

File renamed without changes.

0 commit comments

Comments
 (0)