feat: merge the last thinking with the previous group

This commit is contained in:
Henry Li
2026-01-17 16:10:58 +08:00
parent be1e016ed4
commit 1a3b70ac43
2 changed files with 15 additions and 22 deletions

View File

@@ -70,13 +70,6 @@ function MessageContent_({
const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading); const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading);
return ( return (
<AIElementMessageContent className={className}> <AIElementMessageContent className={className}>
{hasReasoning(message) && (
<MessageGroup
className="mb-2"
messages={[message]}
isLoading={isLoading}
/>
)}
<AIElementMessageResponse rehypePlugins={rehypePlugins}> <AIElementMessageResponse rehypePlugins={rehypePlugins}>
{extractContentFromMessage(message)} {extractContentFromMessage(message)}
</AIElementMessageResponse> </AIElementMessageResponse>

View File

@@ -2,29 +2,17 @@ import type { Message } from "@langchain/langgraph-sdk";
export function groupMessages<T>( export function groupMessages<T>(
messages: Message[], messages: Message[],
mapper: ( mapper: (groupedMessages: Message[]) => T,
groupedMessages: Message[],
groupIndex: number,
isLastGroup: boolean,
) => T,
isLoading = false, isLoading = false,
): T[] { ): T[] {
if (messages.length === 0) { if (messages.length === 0) {
return []; return [];
} }
const resultsOfGroups: T[] = []; const groups: Message[][] = [];
let currentGroup: Message[] = []; let currentGroup: Message[] = [];
const lastMessage = messages[messages.length - 1]!;
const yieldCurrentGroup = () => { const yieldCurrentGroup = () => {
if (currentGroup.length > 0) { if (currentGroup.length > 0) {
const resultOfGroup = mapper( groups.push(currentGroup);
currentGroup,
resultsOfGroups.length,
currentGroup.includes(lastMessage),
);
if (resultOfGroup !== undefined && resultOfGroup !== null) {
resultsOfGroups.push(resultOfGroup);
}
currentGroup = []; currentGroup = [];
} }
}; };
@@ -47,6 +35,7 @@ export function groupMessages<T>(
isLoading) isLoading)
) { ) {
if (message.tool_calls?.[0]?.name === "present_files") { if (message.tool_calls?.[0]?.name === "present_files") {
// When `present_files` called, put them into an individual group
yieldCurrentGroup(); yieldCurrentGroup();
currentGroup.push(message); currentGroup.push(message);
} else { } else {
@@ -57,6 +46,9 @@ export function groupMessages<T>(
} else { } else {
// Assistant messages with content (text or images) are shown as a group if they have content // Assistant messages with content (text or images) are shown as a group if they have content
// No matter whether it has tool calls or not // No matter whether it has tool calls or not
if (hasReasoning(message)) {
currentGroup.push(message);
}
yieldCurrentGroup(); yieldCurrentGroup();
currentGroup.push(message); currentGroup.push(message);
} }
@@ -64,6 +56,14 @@ export function groupMessages<T>(
messageIndex++; messageIndex++;
} }
yieldCurrentGroup(); yieldCurrentGroup();
const resultsOfGroups: T[] = [];
for (const group of groups) {
const resultOfGroup = mapper(group);
if (resultOfGroup !== undefined && resultOfGroup !== null) {
resultsOfGroups.push(resultOfGroup);
}
}
return resultsOfGroups; return resultsOfGroups;
} }