mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-04-27 17:56:25 +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
|
|
@ -8,20 +8,15 @@ import {
|
|||
import { useAtomValue } from "jotai";
|
||||
import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react";
|
||||
import type { FC } from "react";
|
||||
import { useContext, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { useEffect, useMemo, useRef, useState } from "react";
|
||||
import { commentsEnabledAtom, targetCommentIdAtom } from "@/atoms/chat/current-thread.atom";
|
||||
import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
|
||||
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
|
||||
import {
|
||||
ThinkingStepsContext,
|
||||
ThinkingStepsDisplay,
|
||||
} from "@/components/assistant-ui/thinking-steps";
|
||||
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
|
||||
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
|
||||
import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container";
|
||||
import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet";
|
||||
import { CreateConfluencePageToolUI, DeleteConfluencePageToolUI, UpdateConfluencePageToolUI } from "@/components/tool-ui/confluence";
|
||||
import { DeepAgentThinkingToolUI } from "@/components/tool-ui/deepagent-thinking";
|
||||
import { DisplayImageToolUI } from "@/components/tool-ui/display-image";
|
||||
import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
|
||||
import { GenerateReportToolUI } from "@/components/tool-ui/generate-report";
|
||||
|
|
@ -50,44 +45,15 @@ export const MessageError: FC = () => {
|
|||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Custom component to render thinking steps from Context
|
||||
*/
|
||||
const ThinkingStepsPart: FC = () => {
|
||||
const thinkingStepsMap = useContext(ThinkingStepsContext);
|
||||
|
||||
// Get the current message ID to look up thinking steps
|
||||
const messageId = useAuiState(({ message }) => message?.id);
|
||||
const thinkingSteps = thinkingStepsMap.get(messageId) || [];
|
||||
|
||||
// Check if this specific message is currently streaming
|
||||
// A message is streaming if: thread is running AND this is the last assistant message
|
||||
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
|
||||
const isMessageStreaming = isThreadRunning && isLastMessage;
|
||||
|
||||
if (thinkingSteps.length === 0) return null;
|
||||
|
||||
return (
|
||||
<div className="mb-3">
|
||||
<ThinkingStepsDisplay steps={thinkingSteps} isThreadRunning={isMessageStreaming} />
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const AssistantMessageInner: FC = () => {
|
||||
return (
|
||||
<>
|
||||
{/* Render thinking steps from message content - this ensures proper scroll tracking */}
|
||||
<ThinkingStepsPart />
|
||||
|
||||
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
|
||||
<MessagePrimitive.Parts
|
||||
components={{
|
||||
Text: MarkdownText,
|
||||
tools: {
|
||||
by_name: {
|
||||
deepagent_thinking: DeepAgentThinkingToolUI,
|
||||
generate_report: GenerateReportToolUI,
|
||||
generate_podcast: GeneratePodcastToolUI,
|
||||
generate_video_presentation: GenerateVideoPresentationToolUI,
|
||||
|
|
|
|||
|
|
@ -1,14 +1,12 @@
|
|||
import { makeAssistantDataUI, useAuiState } from "@assistant-ui/react";
|
||||
import { ChevronRightIcon } from "lucide-react";
|
||||
import type { FC } from "react";
|
||||
import { createContext, useCallback, useEffect, useState } from "react";
|
||||
import { useCallback, useEffect, useState } from "react";
|
||||
import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought";
|
||||
import { TextShimmerLoader } from "@/components/prompt-kit/loader";
|
||||
import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
// Context to pass thinking steps to AssistantMessage
|
||||
export const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
|
||||
|
||||
/**
|
||||
* Chain of thought display component - single collapsible dropdown design
|
||||
*/
|
||||
|
|
@ -18,7 +16,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
}) => {
|
||||
const [isOpen, setIsOpen] = useState(true);
|
||||
|
||||
// Derive effective status for each step
|
||||
const getEffectiveStatus = useCallback(
|
||||
(step: ThinkingStep): "pending" | "in_progress" | "completed" => {
|
||||
if (step.status === "in_progress" && !isThreadRunning) {
|
||||
|
|
@ -36,7 +33,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
steps.every((s) => getEffectiveStatus(s) === "completed");
|
||||
const isProcessing = isThreadRunning && !allCompleted;
|
||||
|
||||
// Auto-collapse when all tasks are completed
|
||||
useEffect(() => {
|
||||
if (allCompleted) {
|
||||
setIsOpen(false);
|
||||
|
|
@ -61,7 +57,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
return (
|
||||
<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
|
||||
<div className="rounded-lg">
|
||||
{/* Main collapsible header */}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setIsOpen(!isOpen)}
|
||||
|
|
@ -70,20 +65,17 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
"text-muted-foreground hover:text-foreground"
|
||||
)}
|
||||
>
|
||||
{/* Header text with shimmer if processing (streaming) */}
|
||||
{isProcessing ? (
|
||||
<TextShimmerLoader text={getHeaderText()} size="sm" />
|
||||
) : (
|
||||
<span>{getHeaderText()}</span>
|
||||
)}
|
||||
|
||||
{/* Chevron */}
|
||||
<ChevronRightIcon
|
||||
className={cn("size-4 transition-transform duration-200", isOpen && "rotate-90")}
|
||||
/>
|
||||
</button>
|
||||
|
||||
{/* Collapsible content with CSS grid animation */}
|
||||
<div
|
||||
className={cn(
|
||||
"grid transition-[grid-template-rows] duration-300 ease-out",
|
||||
|
|
@ -98,13 +90,10 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
|
||||
return (
|
||||
<div key={step.id} className="relative flex gap-3">
|
||||
{/* Dot and line column */}
|
||||
<div className="relative flex flex-col items-center w-2">
|
||||
{/* Vertical connection line - extends to next dot */}
|
||||
{!isLast && (
|
||||
<div className="absolute left-1/2 top-[15px] -bottom-[7px] w-px -translate-x-1/2 bg-muted-foreground/30" />
|
||||
)}
|
||||
{/* Step dot - on top of line */}
|
||||
<div className="relative z-10 mt-[7px] flex shrink-0 items-center justify-center">
|
||||
{effectiveStatus === "in_progress" ? (
|
||||
<span className="relative flex size-2">
|
||||
|
|
@ -117,9 +106,7 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
</div>
|
||||
</div>
|
||||
|
||||
{/* Step content */}
|
||||
<div className="flex-1 min-w-0 pb-4">
|
||||
{/* Step title */}
|
||||
<div
|
||||
className={cn(
|
||||
"text-sm leading-5",
|
||||
|
|
@ -131,7 +118,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
{step.title}
|
||||
</div>
|
||||
|
||||
{/* Step items (sub-content) */}
|
||||
{step.items && step.items.length > 0 && (
|
||||
<div className="mt-1 space-y-0.5">
|
||||
{step.items.map((item, idx) => (
|
||||
|
|
@ -153,3 +139,28 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
|
|||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* assistant-ui data UI component that renders thinking steps from message content.
|
||||
* Registered globally via makeAssistantDataUI — renders inside MessagePrimitive.Parts
|
||||
* at the position of the data part in the content array.
|
||||
*/
|
||||
function ThinkingStepsDataRenderer({ data }: { name: string; data: unknown }) {
|
||||
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
|
||||
const isMessageStreaming = isThreadRunning && isLastMessage;
|
||||
|
||||
const steps = (data as { steps: ThinkingStep[] } | null)?.steps ?? [];
|
||||
if (steps.length === 0) return null;
|
||||
|
||||
return (
|
||||
<div className="mb-3 -mx-2 leading-normal">
|
||||
<ThinkingStepsDisplay steps={steps} isThreadRunning={isMessageStreaming} />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export const ThinkingStepsDataUI = makeAssistantDataUI({
|
||||
name: "thinking-steps",
|
||||
render: ThinkingStepsDataRenderer,
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,6 @@
|
|||
import {
|
||||
ActionBarPrimitive,
|
||||
AuiIf,
|
||||
BranchPickerPrimitive,
|
||||
ComposerPrimitive,
|
||||
ErrorPrimitive,
|
||||
MessagePrimitive,
|
||||
ThreadPrimitive,
|
||||
useAui,
|
||||
|
|
@ -14,14 +11,8 @@ import {
|
|||
AlertCircle,
|
||||
ArrowDownIcon,
|
||||
ArrowUpIcon,
|
||||
CheckIcon,
|
||||
ChevronLeftIcon,
|
||||
ChevronRightIcon,
|
||||
CopyIcon,
|
||||
DownloadIcon,
|
||||
Globe,
|
||||
Plus,
|
||||
RefreshCwIcon,
|
||||
Settings2,
|
||||
SquareIcon,
|
||||
Unplug,
|
||||
|
|
@ -32,7 +23,7 @@ import {
|
|||
import { AnimatePresence, motion } from "motion/react";
|
||||
import Image from "next/image";
|
||||
import { useParams } from "next/navigation";
|
||||
import { type FC, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { type FC, useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { createPortal } from "react-dom";
|
||||
import {
|
||||
agentToolsAtom,
|
||||
|
|
@ -63,12 +54,6 @@ import {
|
|||
InlineMentionEditor,
|
||||
type InlineMentionEditorRef,
|
||||
} from "@/components/assistant-ui/inline-mention-editor";
|
||||
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
|
||||
import {
|
||||
ThinkingStepsContext,
|
||||
ThinkingStepsDisplay,
|
||||
} from "@/components/assistant-ui/thinking-steps";
|
||||
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
|
||||
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
|
||||
import { UserMessage } from "@/components/assistant-ui/user-message";
|
||||
import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/components/layout/ui/sidebar/SidebarSlideOutPanel";
|
||||
|
|
@ -76,7 +61,6 @@ import {
|
|||
DocumentMentionPicker,
|
||||
type DocumentMentionPickerRef,
|
||||
} from "@/components/new-chat/document-mention-picker";
|
||||
import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
|
||||
import { Avatar, AvatarFallback, AvatarGroup } from "@/components/ui/avatar";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Drawer, DrawerContent, DrawerHandle, DrawerTitle } from "@/components/ui/drawer";
|
||||
|
|
@ -111,16 +95,8 @@ const CYCLING_PLACEHOLDERS = [
|
|||
"Check if this week's Slack messages reference any GitHub issues",
|
||||
];
|
||||
|
||||
interface ThreadProps {
|
||||
messageThinkingSteps?: Map<string, ThinkingStep[]>;
|
||||
}
|
||||
|
||||
export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map() }) => {
|
||||
return (
|
||||
<ThinkingStepsContext.Provider value={messageThinkingSteps}>
|
||||
<ThreadContent />
|
||||
</ThinkingStepsContext.Provider>
|
||||
);
|
||||
export const Thread: FC = () => {
|
||||
return <ThreadContent />;
|
||||
};
|
||||
|
||||
const ThreadContent: FC = () => {
|
||||
|
|
@ -1132,97 +1108,6 @@ const TOOL_GROUPS: ToolGroup[] = [
|
|||
},
|
||||
];
|
||||
|
||||
const MessageError: FC = () => {
|
||||
return (
|
||||
<MessagePrimitive.Error>
|
||||
<ErrorPrimitive.Root className="aui-message-error-root mt-2 rounded-md border border-destructive bg-destructive/10 p-3 text-destructive text-sm dark:bg-destructive/5 dark:text-red-200">
|
||||
<ErrorPrimitive.Message className="aui-message-error-message line-clamp-2" />
|
||||
</ErrorPrimitive.Root>
|
||||
</MessagePrimitive.Error>
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Custom component to render thinking steps from Context
|
||||
*/
|
||||
const ThinkingStepsPart: FC = () => {
|
||||
const thinkingStepsMap = useContext(ThinkingStepsContext);
|
||||
|
||||
// Get the current message ID to look up thinking steps
|
||||
const messageId = useAuiState(({ message }) => message?.id);
|
||||
const thinkingSteps = thinkingStepsMap.get(messageId) || [];
|
||||
|
||||
// Check if this specific message is currently streaming
|
||||
// A message is streaming if: thread is running AND this is the last assistant message
|
||||
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false);
|
||||
const isMessageStreaming = isThreadRunning && isLastMessage;
|
||||
|
||||
if (thinkingSteps.length === 0) return null;
|
||||
|
||||
return (
|
||||
<div className="mb-3">
|
||||
<ThinkingStepsDisplay steps={thinkingSteps} isThreadRunning={isMessageStreaming} />
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const AssistantMessageInner: FC = () => {
|
||||
return (
|
||||
<>
|
||||
{/* Render thinking steps from message content - this ensures proper scroll tracking */}
|
||||
<ThinkingStepsPart />
|
||||
|
||||
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
|
||||
<MessagePrimitive.Parts
|
||||
components={{
|
||||
Text: MarkdownText,
|
||||
tools: { Fallback: ToolFallback },
|
||||
}}
|
||||
/>
|
||||
<MessageError />
|
||||
</div>
|
||||
|
||||
<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
|
||||
<BranchPicker />
|
||||
<AssistantActionBar />
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
||||
const AssistantActionBar: FC = () => {
|
||||
return (
|
||||
<ActionBarPrimitive.Root
|
||||
hideWhenRunning
|
||||
autohide="not-last"
|
||||
autohideFloat="single-branch"
|
||||
className="aui-assistant-action-bar-root -ml-1 col-start-3 row-start-2 flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
|
||||
>
|
||||
<ActionBarPrimitive.Copy asChild>
|
||||
<TooltipIconButton tooltip="Copy">
|
||||
<AuiIf condition={({ message }) => message.isCopied}>
|
||||
<CheckIcon />
|
||||
</AuiIf>
|
||||
<AuiIf condition={({ message }) => !message.isCopied}>
|
||||
<CopyIcon />
|
||||
</AuiIf>
|
||||
</TooltipIconButton>
|
||||
</ActionBarPrimitive.Copy>
|
||||
<ActionBarPrimitive.ExportMarkdown asChild>
|
||||
<TooltipIconButton tooltip="Export as Markdown">
|
||||
<DownloadIcon />
|
||||
</TooltipIconButton>
|
||||
</ActionBarPrimitive.ExportMarkdown>
|
||||
<ActionBarPrimitive.Reload asChild>
|
||||
<TooltipIconButton tooltip="Refresh">
|
||||
<RefreshCwIcon />
|
||||
</TooltipIconButton>
|
||||
</ActionBarPrimitive.Reload>
|
||||
</ActionBarPrimitive.Root>
|
||||
);
|
||||
};
|
||||
|
||||
const EditComposer: FC = () => {
|
||||
return (
|
||||
<MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
|
||||
|
|
@ -1245,30 +1130,3 @@ const EditComposer: FC = () => {
|
|||
</MessagePrimitive.Root>
|
||||
);
|
||||
};
|
||||
|
||||
const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({ className, ...rest }) => {
|
||||
return (
|
||||
<BranchPickerPrimitive.Root
|
||||
hideWhenSingleBranch
|
||||
className={cn(
|
||||
"aui-branch-picker-root -ml-2 mr-2 inline-flex items-center text-muted-foreground text-xs",
|
||||
className
|
||||
)}
|
||||
{...rest}
|
||||
>
|
||||
<BranchPickerPrimitive.Previous asChild>
|
||||
<TooltipIconButton tooltip="Previous">
|
||||
<ChevronLeftIcon />
|
||||
</TooltipIconButton>
|
||||
</BranchPickerPrimitive.Previous>
|
||||
<span className="aui-branch-picker-state font-medium">
|
||||
<BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
|
||||
</span>
|
||||
<BranchPickerPrimitive.Next asChild>
|
||||
<TooltipIconButton tooltip="Next">
|
||||
<ChevronRightIcon />
|
||||
</TooltipIconButton>
|
||||
</BranchPickerPrimitive.Next>
|
||||
</BranchPickerPrimitive.Root>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue