mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-10 16:22:38 +02:00
refactor: migrate thinking steps handling to new data structure and streamline related components
This commit is contained in:
Anish Sarkar
parent
b8f3f41326
commit
e587b588c9
7 changed files with 135 additions and 353 deletions
|
|
@ -34,10 +34,10 @@ import { closeEditorPanelAtom } from "@/atoms/editor/editor-panel.atom";
|
|||
import { membersAtom } from "@/atoms/members/members-query.atoms";
|
||||
import { currentUserAtom } from "@/atoms/user/user-query.atoms";
|
||||
import { Thread } from "@/components/assistant-ui/thread";
|
||||
import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps";
|
||||
import { MobileEditorPanel } from "@/components/editor-panel/editor-panel";
|
||||
import { MobileHitlEditPanel } from "@/components/hitl-edit-panel/hitl-edit-panel";
|
||||
import { MobileReportPanel } from "@/components/report-panel/report-panel";
|
||||
import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { useChatSessionStateSync } from "@/hooks/use-chat-session-state";
|
||||
import { useMessagesElectric } from "@/hooks/use-messages-electric";
|
||||
|
|
@ -57,6 +57,7 @@ import {
|
|||
type ContentPartsState,
|
||||
readSSEStream,
|
||||
type ThinkingStepData,
|
||||
updateThinkingSteps,
|
||||
updateToolCall,
|
||||
} from "@/lib/chat/streaming-state";
|
||||
import {
|
||||
|
|
@ -93,23 +94,6 @@ function markInterruptsCompleted(contentParts: Array<{ type: string; result?: un
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract thinking steps from message content
|
||||
*/
|
||||
function extractThinkingSteps(content: unknown): ThinkingStep[] {
|
||||
if (!Array.isArray(content)) return [];
|
||||
|
||||
const thinkingPart = content.find(
|
||||
(part: unknown) =>
|
||||
typeof part === "object" &&
|
||||
part !== null &&
|
||||
"type" in part &&
|
||||
(part as { type: string }).type === "thinking-steps"
|
||||
) as { type: "thinking-steps"; steps: ThinkingStep[] } | undefined;
|
||||
|
||||
return thinkingPart?.steps || [];
|
||||
}
|
||||
|
||||
/**
|
||||
* Zod schema for mentioned document info (for type-safe parsing)
|
||||
*/
|
||||
|
|
@ -183,11 +167,6 @@ export default function NewChatPage() {
|
|||
const [currentThread, setCurrentThread] = useState<ThreadRecord | null>(null);
|
||||
const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
|
||||
const [isRunning, setIsRunning] = useState(false);
|
||||
// Store thinking steps per message ID - kept separate from content to avoid
|
||||
// "unsupported part type" errors from assistant-ui
|
||||
const [messageThinkingSteps, setMessageThinkingSteps] = useState<Map<string, ThinkingStep[]>>(
|
||||
new Map()
|
||||
);
|
||||
const abortControllerRef = useRef<AbortController | null>(null);
|
||||
const [pendingInterrupt, setPendingInterrupt] = useState<{
|
||||
threadId: number;
|
||||
|
|
@ -295,7 +274,6 @@ export default function NewChatPage() {
|
|||
setMessages([]);
|
||||
setThreadId(null);
|
||||
setCurrentThread(null);
|
||||
setMessageThinkingSteps(new Map());
|
||||
setMentionedDocuments([]);
|
||||
setSidebarDocuments([]);
|
||||
setMessageDocumentsMap({});
|
||||
|
|
@ -320,18 +298,8 @@ export default function NewChatPage() {
|
|||
const loadedMessages = messagesResponse.messages.map(convertToThreadMessage);
|
||||
setMessages(loadedMessages);
|
||||
|
||||
// Extract and restore thinking steps from persisted messages
|
||||
const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
|
||||
// Extract and restore mentioned documents from persisted messages
|
||||
const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
|
||||
|
||||
for (const msg of messagesResponse.messages) {
|
||||
if (msg.role === "assistant") {
|
||||
const steps = extractThinkingSteps(msg.content);
|
||||
if (steps.length > 0) {
|
||||
restoredThinkingSteps.set(`msg-${msg.id}`, steps);
|
||||
}
|
||||
}
|
||||
if (msg.role === "user") {
|
||||
const docs = extractMentionedDocuments(msg.content);
|
||||
if (docs.length > 0) {
|
||||
|
|
@ -339,9 +307,6 @@ export default function NewChatPage() {
|
|||
}
|
||||
}
|
||||
}
|
||||
if (restoredThinkingSteps.size > 0) {
|
||||
setMessageThinkingSteps(restoredThinkingSteps);
|
||||
}
|
||||
if (Object.keys(restoredDocsMap).length > 0) {
|
||||
setMessageDocumentsMap(restoredDocsMap);
|
||||
}
|
||||
|
|
@ -745,18 +710,17 @@ export default function NewChatPage() {
|
|||
}
|
||||
|
||||
case "data-thinking-step": {
|
||||
// Handle thinking step events for chain-of-thought display
|
||||
const stepData = parsed.data as ThinkingStepData;
|
||||
if (stepData?.id) {
|
||||
currentThinkingSteps.set(stepData.id, stepData);
|
||||
// Update thinking steps state for rendering
|
||||
// The ThinkingStepsScrollHandler in Thread component
|
||||
// will handle auto-scrolling when this state changes
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const newMap = new Map(prev);
|
||||
newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values()));
|
||||
return newMap;
|
||||
});
|
||||
updateThinkingSteps(contentPartsState, currentThinkingSteps);
|
||||
setMessages((prev) =>
|
||||
prev.map((m) =>
|
||||
m.id === assistantMsgId
|
||||
? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
|
||||
: m
|
||||
)
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
|
@ -821,13 +785,8 @@ export default function NewChatPage() {
|
|||
}
|
||||
}
|
||||
|
||||
// Persist assistant message (with thinking steps for restoration on refresh)
|
||||
// Skip persistence for interrupted messages -- handleResume will persist the final version
|
||||
const finalContent = buildContentForPersistence(
|
||||
contentPartsState,
|
||||
TOOLS_WITH_UI,
|
||||
currentThinkingSteps
|
||||
);
|
||||
const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
|
||||
if (contentParts.length > 0 && !wasInterrupted) {
|
||||
try {
|
||||
const savedMessage = await appendMessage(currentThreadId, {
|
||||
|
|
@ -847,18 +806,6 @@ export default function NewChatPage() {
|
|||
? { ...prev, assistantMsgId: newMsgId }
|
||||
: prev
|
||||
);
|
||||
|
||||
// Also update thinking steps map with new ID
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const steps = prev.get(assistantMsgId);
|
||||
if (steps) {
|
||||
const newMap = new Map(prev);
|
||||
newMap.delete(assistantMsgId);
|
||||
newMap.set(newMsgId, steps);
|
||||
return newMap;
|
||||
}
|
||||
return prev;
|
||||
});
|
||||
} catch (err) {
|
||||
console.error("Failed to persist assistant message:", err);
|
||||
}
|
||||
|
|
@ -875,11 +822,7 @@ export default function NewChatPage() {
|
|||
(part.type === "tool-call" && TOOLS_WITH_UI.has(part.toolName))
|
||||
);
|
||||
if (hasContent && currentThreadId) {
|
||||
const partialContent = buildContentForPersistence(
|
||||
contentPartsState,
|
||||
TOOLS_WITH_UI,
|
||||
currentThinkingSteps
|
||||
);
|
||||
const partialContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
|
||||
try {
|
||||
const savedMessage = await appendMessage(currentThreadId, {
|
||||
role: "assistant",
|
||||
|
|
@ -926,7 +869,6 @@ export default function NewChatPage() {
|
|||
} finally {
|
||||
setIsRunning(false);
|
||||
abortControllerRef.current = null;
|
||||
// Note: We no longer clear thinking steps - they persist with the message
|
||||
}
|
||||
},
|
||||
[
|
||||
|
|
@ -969,9 +911,7 @@ export default function NewChatPage() {
|
|||
const controller = new AbortController();
|
||||
abortControllerRef.current = controller;
|
||||
|
||||
const currentThinkingSteps = new Map<string, ThinkingStepData>(
|
||||
(messageThinkingSteps.get(assistantMsgId) ?? []).map((s) => [s.id, s])
|
||||
);
|
||||
const currentThinkingSteps = new Map<string, ThinkingStepData>();
|
||||
|
||||
const contentPartsState: ContentPartsState = {
|
||||
contentParts: [],
|
||||
|
|
@ -998,6 +938,15 @@ export default function NewChatPage() {
|
|||
result: p.result as unknown,
|
||||
});
|
||||
contentPartsState.currentTextPartIndex = -1;
|
||||
} else if (p.type === "data-thinking-steps") {
|
||||
const stepsData = p.data as { steps: ThinkingStepData[] } | undefined;
|
||||
contentParts.push({
|
||||
type: "data-thinking-steps",
|
||||
data: { steps: stepsData?.steps ?? [] },
|
||||
});
|
||||
for (const step of stepsData?.steps ?? []) {
|
||||
currentThinkingSteps.set(step.id, step);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1115,11 +1064,14 @@ export default function NewChatPage() {
|
|||
const stepData = parsed.data as ThinkingStepData;
|
||||
if (stepData?.id) {
|
||||
currentThinkingSteps.set(stepData.id, stepData);
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const newMap = new Map(prev);
|
||||
newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values()));
|
||||
return newMap;
|
||||
});
|
||||
updateThinkingSteps(contentPartsState, currentThinkingSteps);
|
||||
setMessages((prev) =>
|
||||
prev.map((m) =>
|
||||
m.id === assistantMsgId
|
||||
? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
|
||||
: m
|
||||
)
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
|
@ -1173,11 +1125,7 @@ export default function NewChatPage() {
|
|||
}
|
||||
}
|
||||
|
||||
const finalContent = buildContentForPersistence(
|
||||
contentPartsState,
|
||||
TOOLS_WITH_UI,
|
||||
currentThinkingSteps
|
||||
);
|
||||
const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
|
||||
if (contentParts.length > 0) {
|
||||
try {
|
||||
const savedMessage = await appendMessage(resumeThreadId, {
|
||||
|
|
@ -1188,16 +1136,6 @@ export default function NewChatPage() {
|
|||
setMessages((prev) =>
|
||||
prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m))
|
||||
);
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const steps = prev.get(assistantMsgId);
|
||||
if (steps) {
|
||||
const newMap = new Map(prev);
|
||||
newMap.delete(assistantMsgId);
|
||||
newMap.set(newMsgId, steps);
|
||||
return newMap;
|
||||
}
|
||||
return prev;
|
||||
});
|
||||
} catch (err) {
|
||||
console.error("Failed to persist resumed assistant message:", err);
|
||||
}
|
||||
|
|
@ -1213,7 +1151,7 @@ export default function NewChatPage() {
|
|||
abortControllerRef.current = null;
|
||||
}
|
||||
},
|
||||
[pendingInterrupt, messages, searchSpaceId, messageThinkingSteps]
|
||||
[pendingInterrupt, messages, searchSpaceId]
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -1332,20 +1270,6 @@ export default function NewChatPage() {
|
|||
return prev;
|
||||
});
|
||||
|
||||
// Clear thinking steps for the removed messages
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const newMap = new Map(prev);
|
||||
// Remove thinking steps for the last two messages
|
||||
const lastTwoIds = messages
|
||||
.slice(-2)
|
||||
.map((m) => m.id)
|
||||
.filter((id): id is string => !!id);
|
||||
for (const id of lastTwoIds) {
|
||||
newMap.delete(id);
|
||||
}
|
||||
return newMap;
|
||||
});
|
||||
|
||||
// Start streaming
|
||||
setIsRunning(true);
|
||||
const controller = new AbortController();
|
||||
|
|
@ -1476,11 +1400,14 @@ export default function NewChatPage() {
|
|||
const stepData = parsed.data as ThinkingStepData;
|
||||
if (stepData?.id) {
|
||||
currentThinkingSteps.set(stepData.id, stepData);
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const newMap = new Map(prev);
|
||||
newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values()));
|
||||
return newMap;
|
||||
});
|
||||
updateThinkingSteps(contentPartsState, currentThinkingSteps);
|
||||
setMessages((prev) =>
|
||||
prev.map((m) =>
|
||||
m.id === assistantMsgId
|
||||
? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
|
||||
: m
|
||||
)
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
|
@ -1491,11 +1418,7 @@ export default function NewChatPage() {
|
|||
}
|
||||
|
||||
// Persist messages after streaming completes
|
||||
const finalContent = buildContentForPersistence(
|
||||
contentPartsState,
|
||||
TOOLS_WITH_UI,
|
||||
currentThinkingSteps
|
||||
);
|
||||
const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
|
||||
if (contentParts.length > 0) {
|
||||
try {
|
||||
// Persist user message (for both edit and reload modes, since backend deleted it)
|
||||
|
|
@ -1526,18 +1449,6 @@ export default function NewChatPage() {
|
|||
prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m))
|
||||
);
|
||||
|
||||
setMessageThinkingSteps((prev) => {
|
||||
const steps = prev.get(assistantMsgId);
|
||||
if (steps) {
|
||||
const newMap = new Map(prev);
|
||||
newMap.delete(assistantMsgId);
|
||||
newMap.set(newMsgId, steps);
|
||||
return newMap;
|
||||
}
|
||||
return prev;
|
||||
});
|
||||
|
||||
// Track successful response
|
||||
trackChatResponseReceived(searchSpaceId, threadId);
|
||||
} catch (err) {
|
||||
console.error("Failed to persist regenerated message:", err);
|
||||
|
|
@ -1570,7 +1481,7 @@ export default function NewChatPage() {
|
|||
abortControllerRef.current = null;
|
||||
}
|
||||
},
|
||||
[threadId, searchSpaceId, messages, setMessageThinkingSteps, disabledTools]
|
||||
[threadId, searchSpaceId, messages, disabledTools]
|
||||
);
|
||||
|
||||
// Handle editing a message - truncates history and regenerates with new query
|
||||
|
|
@ -1675,9 +1586,10 @@ export default function NewChatPage() {
|
|||
|
||||
return (
|
||||
<AssistantRuntimeProvider runtime={runtime}>
|
||||
<ThinkingStepsDataUI />
|
||||
<div key={searchSpaceId} className="flex h-[calc(100dvh-64px)] overflow-hidden">
|
||||
<div className="flex-1 flex flex-col min-w-0 overflow-hidden">
|
||||
<Thread messageThinkingSteps={messageThinkingSteps} />
|
||||
<Thread />
|
||||
</div>
|
||||
<MobileReportPanel />
|
||||
<MobileEditorPanel />
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue