From 304755aba3d2b9f87358ade56cfafbe44253c8dc Mon Sep 17 00:00:00 2001 From: Innei Date: Tue, 23 Sep 2025 20:46:45 +0800 Subject: [PATCH] fix: refactor AIChainOfThought for improved collapse functionality - Introduced a new `AIInnerReasoningPart` component to encapsulate the rendering logic for individual reasoning parts, enhancing modularity. - Updated `AIChainOfThought` to utilize the new component and manage collapse state more effectively using a ref. - Simplified the extraction of headings from reasoning text, improving readability and maintainability of the code. These changes aim to enhance the user experience by providing clearer feedback during AI reasoning processes and improving the overall structure of the components. Signed-off-by: Innei --- .../components/displays/AIChainOfThought.tsx | 148 ++++++++++-------- .../src/ui/collapse/CollapseCss.tsx | 10 ++ 2 files changed, 93 insertions(+), 65 deletions(-) diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx b/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx index aeb232500..2b80bd327 100644 --- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx +++ b/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx @@ -1,3 +1,4 @@ +import type { CollapseCssRef } from "@follow/components/ui/collapse/CollapseCss.js" import { CollapseCss, CollapseCssGroup } from "@follow/components/ui/collapse/CollapseCss.js" import { cn } from "@follow/utils" import type { ReasoningUIPart } from "ai" @@ -10,62 +11,41 @@ interface AIChainOfThoughtProps { isStreaming?: boolean className?: string } - export const AIChainOfThought: React.FC = React.memo( - ({ groups, isStreaming = false, className }) => { + ({ groups, isStreaming, className }) => { const collapseId = React.useMemo(() => `chain-${Math.random().toString(36).slice(2)}`, []) - // Re-mount CollapseCssGroup when streaming state changes or when we need to force-open while streaming - const [remountTick, setRemountTick] = React.useState(0) - const groupKey = `${isStreaming ? "streaming" : "idle"}:${remountTick}` + const collapseRef = React.useRef(null) const lastPartText = groups.at?.(-1)?.text + const currentChainReasoningIsFinished = React.useMemo(() => { + return groups.every((part) => part.state === "done") + }, [groups]) const currentReasoningTitle = React.useMemo(() => { if (!isStreaming) return null return extractHeading(lastPartText) }, [isStreaming, lastPartText]) + React.useEffect(() => { + collapseRef.current?.setIsOpened(!currentChainReasoningIsFinished) + }, [collapseRef, currentChainReasoningIsFinished]) + if (!groups || groups.length === 0) return null - const extractHeading = (text?: string): string | undefined => { - if (!text) return - const lines = text.split(/\r?\n/) - for (const raw of lines) { - const line = raw.trim() - if (!line) continue - if (line.startsWith("#")) { - let idx = 0 - while (idx < line.length && line.charAt(idx) === "#") idx++ - let content = line.slice(idx).trim() - while (content.endsWith("#")) content = content.slice(0, -1).trim() - return content || undefined - } - if (line.startsWith("**") && line.endsWith("**") && line.length > 4) { - return line.slice(2, -2).trim() || undefined - } - break - } - return - } return (
- + { - // While streaming, keep it open and block manual collapse - if (isStreaming && !opened) { - setRemountTick((x) => x + 1) - } - }} + defaultOpen={!currentChainReasoningIsFinished} title={
- {isStreaming ? ( + {!currentChainReasoningIsFinished ? ( Thinking: {currentReasoningTitle} @@ -98,42 +78,17 @@ export const AIChainOfThought: React.FC = React.memo( aria-hidden className={cn( "absolute left-2 top-2 size-2 -translate-x-1/2 rounded-full border", - groupStreaming ? "border-blue bg-blue" : "border-fill bg-fill-vibrant", + groupStreaming ? "border-accent bg-accent" : "border-fill bg-fill-vibrant", )} >
- { - if (groupStreaming && !opened) { - setRemountTick((x) => x + 1) - } - }} - title={ -
-
- {title ? ( - - {"Reason: "} - {title} - - ) : ( - {groupStreaming ? "Reasoning..." : "Reasoning"} - )} -
-
- -
-
- } - className="group/inner w-full border-none" - > - - +
) })} @@ -145,4 +100,67 @@ export const AIChainOfThought: React.FC = React.memo( }, ) +const AIInnerReasoningPart: React.FC<{ + title: string | undefined + text: string + groupStreaming: boolean +}> = React.memo(({ title, text, groupStreaming }) => { + const id = React.useId() + const collapseRef = React.useRef(null) + + React.useEffect(() => { + collapseRef.current?.setIsOpened(groupStreaming) + }, [groupStreaming, collapseRef]) + + return ( + +
+ {title ? ( + + {"Reason: "} + {title} + + ) : ( + {groupStreaming ? "Reasoning..." : "Reasoning"} + )} +
+
+ +
+
+ } + className="group/inner w-full border-none" + > + + + ) +}) + AIChainOfThought.displayName = "AIChainOfThought" + +const extractHeading = (text?: string): string | undefined => { + if (!text) return + const lines = text.split(/\r?\n/) + for (const raw of lines) { + const line = raw.trim() + if (!line) continue + if (line.startsWith("#")) { + let idx = 0 + while (idx < line.length && line.charAt(idx) === "#") idx++ + let content = line.slice(idx).trim() + while (content.endsWith("#")) content = content.slice(0, -1).trim() + return content || undefined + } + if (line.startsWith("**") && line.endsWith("**") && line.length > 4) { + return line.slice(2, -2).trim() || undefined + } + break + } + return +} diff --git a/packages/internal/components/src/ui/collapse/CollapseCss.tsx b/packages/internal/components/src/ui/collapse/CollapseCss.tsx index f6765e773..d7f0807bc 100644 --- a/packages/internal/components/src/ui/collapse/CollapseCss.tsx +++ b/packages/internal/components/src/ui/collapse/CollapseCss.tsx @@ -70,8 +70,12 @@ interface CollapseProps { className?: string children: React.ReactNode innerClassName?: string + ref?: React.Ref } +export interface CollapseCssRef { + setIsOpened: (isOpened: boolean) => void +} export const CollapseCss: FC = ({ title, hideArrow, @@ -83,6 +87,7 @@ export const CollapseCss: FC = ({ className, innerClassName, children, + ref, }) => { const reactId = React.useId() const id = collapseId ?? reactId @@ -100,6 +105,11 @@ export const CollapseCss: FC = ({ onOpenChange?.(newOpened) }, [id, isOpened, controlledIsOpened, setOpenState, onOpenChange]) + React.useImperativeHandle(ref, () => ({ + setIsOpened: (isOpened: boolean) => { + setOpenState(id, isOpened) + }, + })) return (