Skip to content

Commit 71d5507

Browse files
alashchev17alashchev17
committed
feat: saving usage to assistant message on stream & reasoning mode selection fix
1 parent 1bb22c8 commit 71d5507

File tree

11 files changed

+215
-30
lines changed

11 files changed

+215
-30
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import React, { useCallback } from "react";
22
import { Markdown } from "../Markdown";
33

44
import { Container, Box } from "@radix-ui/themes";
5-
import { ToolCall } from "../../services/refact";
5+
import { ToolCall, Usage } from "../../services/refact";
66
import { ToolContent } from "./ToolsContent";
77
import { fallbackCopying } from "../../utils/fallbackCopying";
88
import { telemetryApi } from "../../services/refact/telemetry";
@@ -12,12 +12,14 @@ import { UsageCounter } from "./UsageCounter";
1212
type ChatInputProps = {
1313
message: string | null;
1414
toolCalls?: ToolCall[] | null;
15+
usage?: Usage | null;
1516
isLast?: boolean;
1617
};
1718

1819
export const AssistantInput: React.FC<ChatInputProps> = ({
1920
message,
2021
toolCalls,
22+
usage,
2123
isLast,
2224
}) => {
2325
const [sendTelemetryEvent] =
@@ -60,7 +62,6 @@ export const AssistantInput: React.FC<ChatInputProps> = ({
6062

6163
return (
6264
<Container position="relative">
63-
<UsageCounter />
6465
{message && (
6566
<Box py="4">
6667
<Markdown canHaveInteractiveElements={true} onCopyClick={handleCopy}>
@@ -69,6 +70,7 @@ export const AssistantInput: React.FC<ChatInputProps> = ({
6970
</Box>
7071
)}
7172
{toolCalls && <ToolContent toolCalls={toolCalls} />}
73+
{usage && <UsageCounter usage={usage} />}
7274
{isLast && <LikeButton />}
7375
</Container>
7476
);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,7 @@ function renderMessages(
189189
key={key}
190190
message={head.content}
191191
toolCalls={head.tool_calls}
192+
usage={head.usage}
192193
isLast={isLast}
193194
/>,
194195
];

‎refact-agent/gui/src/components/ChatContent/UsageCounter/UsageCounter.module.css‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
.usageCounterContainer {
2-
position: absolute;
3-
top: 0;
4-
right: 0;
2+
/* position: absolute; */
3+
/* top: 0;
4+
right: 0; */
5+
margin-left: auto;
56
display: flex;
67
align-items: center;
78
padding: var(--space-2) var(--space-3);

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

Lines changed: 122 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,127 @@
1-
import { Card, Flex, Text } from "@radix-ui/themes";
2-
import styles from "./UsageCounter.module.css";
1+
import React from "react";
2+
import { Card, Flex, HoverCard, Text } from "@radix-ui/themes";
33
import { ArrowDownIcon, ArrowUpIcon } from "@radix-ui/react-icons";
44

5-
export const UsageCounter = () => {
5+
import { ScrollArea } from "../../ScrollArea";
6+
import { Usage } from "../../../services/refact";
7+
8+
import styles from "./UsageCounter.module.css";
9+
10+
type UsageCounterProps = {
11+
usage: Usage;
12+
};
13+
/*
14+
15+
completion_tokens: number;
16+
prompt_tokens: number;
17+
total_tokens: number;
18+
completion_tokens_details: CompletionTokenDetails | null;
19+
prompt_tokens_details: PromptTokenDetails | null;
20+
cache_creation_input_tokens?: number;
21+
cache_read_input_tokens?: number;
22+
23+
*/
24+
25+
function formatNumber(num: number): string {
26+
if (num >= 1000000) {
27+
return (num / 1000000).toFixed(1) + "M";
28+
} else if (num >= 1000) {
29+
return (num / 1000).toFixed(2) + "k";
30+
}
31+
return num.toString();
32+
}
33+
34+
export const UsageCounter: React.FC<UsageCounterProps> = ({ usage }) => {
35+
const inputTokens = Object.entries(usage).reduce((acc, [key, value]) => {
36+
if (key === "prompt_tokens" && typeof value === "number") {
37+
return acc + value;
38+
} else if (
39+
key === "cache_creation_input_tokens" &&
40+
typeof value === "number"
41+
) {
42+
return acc + value;
43+
} else if (key === "cache_read_input_tokens" && typeof value === "number") {
44+
return acc + value;
45+
}
46+
return acc;
47+
}, 0);
48+
49+
const outputTokens = Object.entries(usage).reduce((acc, [key, value]) => {
50+
if (key === "completion_tokens" && typeof value === "number") {
51+
return acc + value;
52+
}
53+
return acc;
54+
}, 0);
55+
656
return (
7-
<Card className={styles.usageCounterContainer}>
8-
<Flex align="center">
9-
<ArrowUpIcon width="12" height="12" />
10-
<Text size="1">1.2k</Text>
11-
</Flex>
12-
<Flex align="center">
13-
<ArrowDownIcon width="12" height="12" />
14-
<Text size="1">12k</Text>
15-
</Flex>
16-
</Card>
57+
<HoverCard.Root>
58+
<HoverCard.Trigger>
59+
<Card className={styles.usageCounterContainer}>
60+
<Flex align="center">
61+
<ArrowUpIcon width="12" height="12" />
62+
<Text size="1">{formatNumber(inputTokens)}</Text>
63+
</Flex>
64+
<Flex align="center">
65+
<ArrowDownIcon width="12" height="12" />
66+
<Text size="1">{outputTokens}</Text>
67+
</Flex>
68+
</Card>
69+
</HoverCard.Trigger>
70+
<ScrollArea scrollbars="both" asChild>
71+
<HoverCard.Content
72+
size="1"
73+
maxHeight="50vh"
74+
maxWidth="90vw"
75+
minWidth="300px"
76+
avoidCollisions
77+
align="end"
78+
side="top"
79+
>
80+
<Flex direction="column" align="start" gap="2">
81+
<Text size="2" mb="2">
82+
Tokens spent per message:
83+
</Text>
84+
<Flex align="center" justify="between" width="100%">
85+
<Text size="1" weight="bold">
86+
Input tokens (in total):{" "}
87+
</Text>
88+
<Text size="1">{inputTokens}</Text>
89+
</Flex>
90+
{usage.cache_read_input_tokens ? (
91+
<Flex align="center" justify="between" width="100%">
92+
<Text size="1" weight="bold">
93+
Cache read input tokens:{" "}
94+
</Text>
95+
<Text size="1">{usage.cache_read_input_tokens}</Text>
96+
</Flex>
97+
) : undefined}
98+
{usage.cache_creation_input_tokens && (
99+
<Flex align="center" justify="between" width="100%">
100+
<Text size="1" weight="bold">
101+
Cache creation input tokens:{" "}
102+
</Text>
103+
<Text size="1">{usage.cache_creation_input_tokens}</Text>
104+
</Flex>
105+
)}
106+
<Flex align="center" justify="between" width="100%">
107+
<Text size="1" weight="bold">
108+
Completion tokens:{" "}
109+
</Text>
110+
<Text size="1">{outputTokens}</Text>
111+
</Flex>
112+
{usage.completion_tokens_details && (
113+
<Flex align="center" justify="between" width="100%">
114+
<Text size="1" weight="bold">
115+
Reasoning tokens:{" "}
116+
</Text>
117+
<Text size="1">
118+
{usage.completion_tokens_details.reasoning_tokens}
119+
</Text>
120+
</Flex>
121+
)}
122+
</Flex>
123+
</HoverCard.Content>
124+
</ScrollArea>
125+
</HoverCard.Root>
17126
);
18127
};

‎refact-agent/gui/src/features/Chat/Thread/reducer.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ const getThreadMode = ({
8282
return maybeMode === "CONFIGURE" ? "AGENT" : maybeMode;
8383
}
8484

85-
return chatModeToLspMode(tool_use);
85+
return chatModeToLspMode({ toolUse: tool_use });
8686
};
8787

8888
const createInitialState = ({
@@ -113,7 +113,7 @@ export const chatReducer = createReducer(initialState, (builder) => {
113113
builder.addCase(setToolUse, (state, action) => {
114114
state.thread.tool_use = action.payload;
115115
state.tool_use = action.payload;
116-
state.thread.mode = chatModeToLspMode(action.payload);
116+
state.thread.mode = chatModeToLspMode({ toolUse: action.payload });
117117
});
118118

119119
builder.addCase(setPreventSend, (state, action) => {

‎refact-agent/gui/src/features/Chat/Thread/types.ts‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,12 +83,33 @@ export type LspChatMode =
8383
| "PROJECT_SUMMARY"
8484
| "THINKING_AGENT";
8585

86-
export function chatModeToLspMode(
86+
export function chatModeToLspMode({
87+
toolUse,
88+
mode,
89+
defaultMode,
90+
}: {
91+
toolUse?: ToolUse;
92+
mode?: LspChatMode;
93+
defaultMode?: LspChatMode;
94+
}): LspChatMode {
95+
if (defaultMode) {
96+
if (defaultMode === "AGENT" || defaultMode === "THINKING_AGENT")
97+
return "AGENT";
98+
return defaultMode;
99+
}
100+
if (mode) {
101+
return mode;
102+
}
103+
if (toolUse === "agent") return "AGENT";
104+
if (toolUse === "quick") return "NO_TOOLS";
105+
return "EXPLORE";
106+
}
107+
108+
export function chatModeToLspModeForChat(
87109
toolUse?: ToolUse,
88110
mode?: LspChatMode,
89111
): LspChatMode {
90112
if (mode) {
91-
if (mode === "AGENT" || mode === "THINKING_AGENT") return "AGENT";
92113
return mode;
93114
}
94115
if (toolUse === "agent") return "AGENT";

‎refact-agent/gui/src/features/Chat/Thread/utils.ts‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -210,9 +210,11 @@ export function formatChatResponse(
210210
return messages;
211211
}
212212

213+
const usage = response.usage;
214+
213215
return response.choices.reduce<ChatMessages>((acc, cur) => {
214216
if (isChatContextFileDelta(cur.delta)) {
215-
const msg = { role: cur.delta.role, content: cur.delta.content };
217+
const msg = { role: cur.delta.role, content: cur.delta.content, usage };
216218
return acc.concat([msg]);
217219
}
218220

@@ -228,6 +230,7 @@ export function formatChatResponse(
228230
content: cur.delta.content,
229231
tool_calls: cur.delta.tool_calls,
230232
finish_reason: cur.finish_reason,
233+
usage: response.usage,
231234
};
232235
return acc.concat([msg]);
233236
}
@@ -236,6 +239,7 @@ export function formatChatResponse(
236239
role: cur.delta.role,
237240
content: cur.delta.content,
238241
finish_reason: cur.finish_reason,
242+
usage: response.usage,
239243
} as ChatMessage;
240244
return acc.concat([message]);
241245
}
@@ -250,6 +254,7 @@ export function formatChatResponse(
250254
content: cur.delta.content ?? "",
251255
tool_calls: cur.delta.tool_calls,
252256
finish_reason: cur.finish_reason,
257+
usage: response.usage,
253258
},
254259
]);
255260
}
@@ -268,6 +273,7 @@ export function formatChatResponse(
268273
content: message,
269274
tool_calls: calls,
270275
finish_reason: cur.finish_reason,
276+
usage: response.usage,
271277
},
272278
]);
273279
}
@@ -286,6 +292,7 @@ export function formatChatResponse(
286292
content: currentMessage + cur.delta.content,
287293
tool_calls: toolCalls,
288294
finish_reason: cur.finish_reason,
295+
usage: response.usage ?? lastMessage.usage,
289296
},
290297
]);
291298
} else if (
@@ -297,15 +304,31 @@ export function formatChatResponse(
297304
role: "assistant",
298305
content: cur.delta.content,
299306
finish_reason: cur.finish_reason,
307+
usage: response.usage,
300308
},
301309
]);
302310
} else if (cur.delta.role === "assistant") {
303311
// empty message from JB
304312
return acc;
305313
}
306314

307-
if (cur.delta.role === null || cur.finish_reason !== null) {
308-
return acc;
315+
// saving usage if is assistant message and no delta role
316+
if (
317+
(cur.delta.role === null || cur.finish_reason !== null) &&
318+
isAssistantMessage(lastMessage)
319+
) {
320+
const last = acc.slice(0, -1);
321+
const currentMessage = lastMessage.content ?? "";
322+
const toolCalls = lastMessage.tool_calls;
323+
return last.concat([
324+
{
325+
role: "assistant",
326+
content: currentMessage,
327+
tool_calls: toolCalls,
328+
finish_reason: cur.finish_reason,
329+
usage: response.usage ?? lastMessage.usage,
330+
},
331+
]);
309332
}
310333

311334
// console.log("Fall though");
@@ -452,7 +475,10 @@ export function formatMessagesForChat(
452475
if (message.role === "assistant") {
453476
// TODO: why type cast this.
454477
const assistantMessage = message as AssistantMessage;
455-
return acc.concat(assistantMessage);
478+
return acc.concat({
479+
...assistantMessage,
480+
usage: assistantMessage.usage,
481+
});
456482
}
457483

458484
if (

‎refact-agent/gui/src/hooks/useLinksFromLsp.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ export function useGetLinksFromLsp() {
7171
chat_id: chatId,
7272
messages,
7373
model: model ?? "",
74-
mode: chatModeToLspMode(undefined, threadMode),
74+
mode: chatModeToLspMode({ defaultMode: threadMode }),
7575
current_config_file: maybeIntegration?.path,
7676
},
7777
{ skip: skipLinksRequest },

‎refact-agent/gui/src/hooks/useSendChatRequest.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ import {
4646
} from "../features/ToolConfirmation/confirmationSlice";
4747
import {
4848
chatModeToLspMode,
49+
chatModeToLspModeForChat,
4950
LspChatMode,
5051
setChatMode,
5152
setIsWaitingForResponse,
@@ -173,7 +174,8 @@ export const useSendChatRequest = () => {
173174
dispatch(backUpMessages({ id: chatId, messages }));
174175
dispatch(chatAskedQuestion({ id: chatId }));
175176

176-
const mode = maybeMode ?? chatModeToLspMode(toolUse, threadMode);
177+
const mode =
178+
maybeMode ?? chatModeToLspMode({ toolUse, mode: threadMode });
177179

178180
const toolsConfirmed =
179181
isCurrentToolCallAPatch && isPatchAutomatic
@@ -262,10 +264,13 @@ export const useSendChatRequest = () => {
262264

263265
// TODO: make a better way for setting / detecting thread mode.
264266
const maybeConfigure = threadIntegration ? "CONFIGURE" : undefined;
265-
const mode = chatModeToLspMode(
267+
const mode = chatModeToLspModeForChat(
266268
toolUse,
267269
maybeMode ?? threadMode ?? maybeConfigure,
268270
);
271+
console.log(
272+
`[DEBUG]: maybeMode: ${maybeMode}, threadMode: ${threadMode}, maybeConfigure: ${maybeConfigure}`,
273+
);
269274
dispatch(setChatMode(mode));
270275

271276
void sendMessages(messages, mode);

0 commit comments

Comments
 (0)