Skip to content

Commit c98baef

Browse files
samolegosamolego
committed
bugfixes for streaming calls
1 parent 3f76cac commit c98baef

File tree

4 files changed

+136
-106
lines changed

4 files changed

+136
-106
lines changed

‎tools/server/webui/src/components/ChatMessage.tsx‎

Lines changed: 126 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { useMemo, useState, Fragment } from 'react';
1+
import { useMemo, useState } from 'react';
22
import { useAppContext } from '../utils/app.context';
33
import { Message, PendingMessage } from '../utils/types';
44
import { classNames } from '../utils/misc';
@@ -20,6 +20,94 @@ interface SplitMessage {
2020
isThinking?: boolean;
2121
}
2222

23+
// Helper function to extract thoughts from message content
24+
function extractThoughts(content: string | null, role: string): SplitMessage {
25+
if (content === null || (role !== 'assistant' && role !== 'tool')) {
26+
return { content };
27+
}
28+
29+
let actualContent = '';
30+
let thought = '';
31+
let isThinking = false;
32+
let thinkSplit = content.split('<think>', 2);
33+
actualContent += thinkSplit[0];
34+
35+
while (thinkSplit[1] !== undefined) {
36+
thinkSplit = thinkSplit[1].split('</think>', 2);
37+
thought += thinkSplit[0];
38+
isThinking = true;
39+
if (thinkSplit[1] !== undefined) {
40+
isThinking = false;
41+
thinkSplit = thinkSplit[1].split('<think>', 2);
42+
actualContent += thinkSplit[0];
43+
}
44+
}
45+
46+
return { content: actualContent, thought, isThinking };
47+
}
48+
49+
// Helper component to render a single message part
50+
function MessagePart({
51+
message,
52+
isPending,
53+
showThoughts = true,
54+
className = '',
55+
baseClassName = '',
56+
isMainMessage = false,
57+
}: {
58+
message: Message | PendingMessage;
59+
isPending?: boolean;
60+
showThoughts?: boolean;
61+
className?: string;
62+
baseClassName?: string;
63+
isMainMessage?: boolean;
64+
}) {
65+
const { config } = useAppContext();
66+
const { content, thought, isThinking } = extractThoughts(message.content, message.role);
67+
68+
if (message.role === 'tool' && baseClassName) {
69+
return (
70+
<ToolCallResultDisplay
71+
content={content || ''}
72+
baseClassName={baseClassName}
73+
/>
74+
);
75+
}
76+
77+
return (
78+
<div className={className}>
79+
{showThoughts && thought && (
80+
<ThoughtProcess
81+
isThinking={!!isThinking && !!isPending}
82+
content={thought}
83+
open={config.showThoughtInProgress}
84+
/>
85+
)}
86+
87+
{message.role === 'tool' && content ? (
88+
<ToolCallResultDisplay content={content} />
89+
) : (
90+
content &&
91+
content.trim() !== '' && (
92+
<MarkdownDisplay
93+
content={content}
94+
isGenerating={isPending}
95+
/>
96+
)
97+
)}
98+
99+
{message.tool_calls &&
100+
message.tool_calls.map((toolCall) => (
101+
<ToolCallArgsDisplay
102+
key={toolCall.id}
103+
toolCall={toolCall}
104+
{...(!isMainMessage && baseClassName ? { baseClassName } : {})}
105+
/>
106+
))}
107+
</div>
108+
);
109+
}
110+
23111
export default function ChatMessage({
24112
msg,
25113
chainedParts,
@@ -59,49 +147,18 @@ export default function ChatMessage({
59147
const nextSibling = siblingLeafNodeIds[siblingCurrIdx + 1];
60148
const prevSibling = siblingLeafNodeIds[siblingCurrIdx - 1];
61149

62-
// for reasoning model, we split the message into content and thought
63-
// TODO: implement this as remark/rehype plugin in the future
64-
const {
65-
content: mainDisplayableContent,
66-
thought,
67-
isThinking,
68-
}: SplitMessage = useMemo(() => {
69-
if (
70-
msg.content === null ||
71-
(msg.role !== 'assistant' && msg.role !== 'tool')
72-
) {
73-
return { content: msg.content };
74-
}
75-
let actualContent = '';
76-
let thought = '';
77-
let isThinking = false;
78-
let thinkSplit = msg.content.split('<think>', 2);
79-
actualContent += thinkSplit[0];
80-
while (thinkSplit[1] !== undefined) {
81-
// <think> tag found
82-
thinkSplit = thinkSplit[1].split('</think>', 2);
83-
thought += thinkSplit[0];
84-
isThinking = true;
85-
if (thinkSplit[1] !== undefined) {
86-
// </think> closing tag found
87-
isThinking = false;
88-
thinkSplit = thinkSplit[1].split('<think>', 2);
89-
actualContent += thinkSplit[0];
90-
}
91-
}
92-
93-
return { content: actualContent, thought, isThinking };
94-
}, [msg]);
150+
const mainSplitMessage = useMemo(() => extractThoughts(msg.content, msg.role), [msg.content, msg.role]);
95151

96152
if (!viewingChat) return null;
97153

98154
const toolCalls = msg.tool_calls ?? null;
99155

100156
const hasContentInMainMsg =
101-
mainDisplayableContent && mainDisplayableContent.trim() !== '';
102-
const hasContentInChainedParts = chainedParts?.some(
103-
(part) => part.content && part.content.trim() !== ''
104-
);
157+
mainSplitMessage.content && mainSplitMessage.content.trim() !== '';
158+
const hasContentInChainedParts = chainedParts?.some((part) => {
159+
const splitPart = extractThoughts(part.content, part.role);
160+
return splitPart.content && splitPart.content.trim() !== '';
161+
});
105162
const entireTurnHasSomeDisplayableContent =
106163
hasContentInMainMsg || hasContentInChainedParts;
107164
const isUser = msg.role === 'user';
@@ -162,7 +219,7 @@ export default function ChatMessage({
162219
{/* not editing content, render message */}
163220
{editingContent === null && (
164221
<>
165-
{mainDisplayableContent === null &&
222+
{mainSplitMessage.content === null &&
166223
!toolCalls &&
167224
!chainedParts?.length ? (
168225
<>
@@ -171,63 +228,29 @@ export default function ChatMessage({
171228
</>
172229
) : (
173230
<>
174-
{/* render message as markdown */}
231+
{/* render main message */}
175232
<div dir="auto" tabIndex={0}>
176-
{thought && (
177-
<ThoughtProcess
178-
isThinking={!!isThinking && !!isPending}
179-
content={thought}
180-
open={config.showThoughtInProgress}
181-
/>
182-
)}
183-
184-
{msg.role === 'tool' && mainDisplayableContent ? (
185-
<ToolCallResultDisplay content={mainDisplayableContent} />
186-
) : (
187-
mainDisplayableContent &&
188-
mainDisplayableContent.trim() !== '' && (
189-
<MarkdownDisplay
190-
content={mainDisplayableContent}
191-
isGenerating={isPending}
192-
/>
193-
)
194-
)}
233+
<MessagePart
234+
message={msg}
235+
isPending={isPending}
236+
showThoughts={true}
237+
isMainMessage={true}
238+
/>
195239
</div>
196-
</>
197-
)}
198-
{toolCalls &&
199-
toolCalls.map((toolCall) => (
200-
<ToolCallArgsDisplay key={toolCall.id} toolCall={toolCall} />
201-
))}
202240

203-
{chainedParts?.map((part) => (
204-
<Fragment key={part.id}>
205-
{part.role === 'tool' && part.content && (
206-
<ToolCallResultDisplay
207-
content={part.content}
208-
baseClassName="collapse bg-base-200 collapse-arrow mb-4 mt-2"
241+
{/* render chained parts */}
242+
{chainedParts?.map((part) => (
243+
<MessagePart
244+
key={part.id}
245+
message={part}
246+
isPending={isPending}
247+
showThoughts={true}
248+
className={part.role === 'assistant' ? 'mt-2' : ''}
249+
baseClassName={part.role === 'tool' ? 'collapse bg-base-200 collapse-arrow mb-4 mt-2' : ''}
209250
/>
210-
)}
211-
212-
{part.role === 'assistant' && part.content && (
213-
<div dir="auto" className="mt-2">
214-
<MarkdownDisplay
215-
content={part.content}
216-
isGenerating={!!isPending}
217-
/>
218-
</div>
219-
)}
220-
221-
{part.tool_calls &&
222-
part.tool_calls.map((toolCall) => (
223-
<ToolCallArgsDisplay
224-
key={toolCall.id}
225-
toolCall={toolCall}
226-
baseClassName="collapse bg-base-200 collapse-arrow mb-4 mt-2"
227-
/>
228-
))}
229-
</Fragment>
230-
))}
251+
))}
252+
</>
253+
)}
231254
{/* render timings if enabled */}
232255
{timings && config.showTokensPerSecond && (
233256
<div className="dropdown dropdown-hover dropdown-top mt-2">
@@ -332,10 +355,16 @@ export default function ChatMessage({
332355
<CopyButton
333356
className="badge btn-mini show-on-hover mr-2"
334357
content={
335-
msg.content ??
336-
chainedParts?.find((p) => p.role === 'assistant' && p.content)
337-
?.content ??
338-
''
358+
[msg, ...(chainedParts || [])]
359+
.filter((p) => p.content)
360+
.map((p) => {
361+
if (p.role === 'user') {
362+
return p.content;
363+
} else {
364+
return extractThoughts(p.content, p.role).content;
365+
}
366+
})
367+
.join('\n\n') || ''
339368
}
340369
/>
341370
)}
@@ -390,4 +419,4 @@ function ThoughtProcess({
390419
</div>
391420
</div>
392421
);
393-
}
422+
}

‎tools/server/webui/src/utils/app.context.tsx‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -358,15 +358,15 @@ export const AppContextProvider = ({
358358
);
359359
}
360360

361+
setPending(convId, null);
362+
onChunk(lastMsgId); // trigger scroll to bottom and switch to the last node
363+
361364
// if message ended due to "finish_reason": "tool_calls"
362365
// resend it to assistant to process the result.
363366
if (shouldContinueChain) {
364367
lastMsgId = await generateMessage(convId, lastMsgId, onChunk);
365368
}
366369

367-
setPending(convId, null);
368-
onChunk(lastMsgId); // trigger scroll to bottom and switch to the last node
369-
370370
// Fetch messages from DB
371371
const savedMsgs = await StorageUtils.getMessages(convId);
372372
console.log({ savedMsgs });
@@ -452,7 +452,7 @@ export const AppContextProvider = ({
452452
if (content !== null) {
453453
const now = Date.now();
454454
const currMsgId = now;
455-
StorageUtils.appendMsg(
455+
await StorageUtils.appendMsg(
456456
{
457457
id: currMsgId,
458458
timestamp: now,

‎tools/server/webui/src/utils/tool_calling/agent_tool.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import StorageUtils from '../storage';
21
import {
32
ToolCallRequest,
43
ToolCallOutput,
@@ -11,7 +10,8 @@ export abstract class AgentTool {
1110
public readonly id: string,
1211
public readonly name: string,
1312
public readonly toolDescription: string,
14-
public readonly parameters: ToolCallParameters
13+
public readonly parameters: ToolCallParameters,
14+
private readonly _enabled: () => boolean
1515
) {}
1616

1717
/**
@@ -43,7 +43,7 @@ export abstract class AgentTool {
4343
* @returns enabled status.
4444
*/
4545
public get enabled(): boolean {
46-
return StorageUtils.getConfig()[`tool_${this.id}_enabled`] ?? false;
46+
return this._enabled();
4747
}
4848

4949
/**

‎tools/server/webui/src/utils/tool_calling/js_repl_tool.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ export class JSReplAgentTool extends AgentTool {
2828
super(
2929
JSReplAgentTool.ID,
3030
'Javascript interpreter',
31-
'Executes JavaScript code in a sandboxed iframe. The code should be self-contained valid javascript. You can use console.log(variable) to print out intermediate values, which will be captured.',
31+
'Executes JavaScript code in a sandboxed iframe. The code should be self-contained valid javascript. Only console.log(variable) and final result are included in response content.',
3232
{
3333
type: 'object',
3434
properties: {
@@ -38,7 +38,8 @@ export class JSReplAgentTool extends AgentTool {
3838
},
3939
},
4040
required: ['code'],
41-
} as ToolCallParameters
41+
} as ToolCallParameters,
42+
() => true
4243
);
4344
this.initIframe();
4445
}

0 commit comments

Comments
 (0)