Просмотр исходного кода

Fix thinking blocks with quotes + add handling `[THINK]...[/THINK]` blocks (#16326)

* fix: prevent reasoning blocks with quotes from being truncated

* chore: update webui build output

* feat: Improve thinking content parsing

* test: Adds ChatMessage component stories for different thinking blocks

* chore: update webui build output

* fix: ChatMessage story fix

---------

Co-authored-by: Aleksander Grygier <aleksander.grygier@gmail.com>
Pascal 3 месяцев назад
Родитель
Сommit
5f7e166cbf

BIN
tools/server/public/index.html.gz


+ 83 - 30
tools/server/webui/src/lib/utils/thinking.ts

@@ -1,7 +1,8 @@
 /**
 /**
- * Parses thinking content from a message that may contain <think> tags
+ * Parses thinking content from a message that may contain <think> tags or [THINK] tags
  * Returns an object with thinking content and cleaned message content
  * Returns an object with thinking content and cleaned message content
- * Handles both complete <think>...</think> blocks and incomplete <think> blocks (streaming)
+ * Handles both complete blocks and incomplete blocks (streaming)
+ * Supports formats: <think>...</think> and [THINK]...[/THINK]
  * @param content - The message content to parse
  * @param content - The message content to parse
  * @returns An object containing the extracted thinking content and the cleaned message content
  * @returns An object containing the extracted thinking content and the cleaned message content
  */
  */
@@ -9,12 +10,11 @@ export function parseThinkingContent(content: string): {
 	thinking: string | null;
 	thinking: string | null;
 	cleanContent: string;
 	cleanContent: string;
 } {
 } {
-	const incompleteMatch = content.includes('<think>') && !content.includes('</think>');
+	const incompleteThinkMatch = content.includes('<think>') && !content.includes('</think>');
+	const incompleteThinkBracketMatch = content.includes('[THINK]') && !content.includes('[/THINK]');
 
 
-	if (incompleteMatch) {
-		// Remove the entire <think>... part from clean content
+	if (incompleteThinkMatch) {
 		const cleanContent = content.split('</think>')?.[1]?.trim();
 		const cleanContent = content.split('</think>')?.[1]?.trim();
-		// Extract everything after <think> as thinking content
 		const thinkingContent = content.split('<think>')?.[1]?.trim();
 		const thinkingContent = content.split('<think>')?.[1]?.trim();
 
 
 		return {
 		return {
@@ -23,12 +23,40 @@ export function parseThinkingContent(content: string): {
 		};
 		};
 	}
 	}
 
 
-	const completeMatch = content.includes('</think>');
+	if (incompleteThinkBracketMatch) {
+		const cleanContent = content.split('[/THINK]')?.[1]?.trim();
+		const thinkingContent = content.split('[THINK]')?.[1]?.trim();
 
 
-	if (completeMatch) {
 		return {
 		return {
-			thinking: content.split('</think>')?.[0]?.trim(),
-			cleanContent: content.split('</think>')?.[1]?.trim()
+			cleanContent,
+			thinking: thinkingContent
+		};
+	}
+
+	const completeThinkMatch = content.match(/<think>([\s\S]*?)<\/think>/);
+	const completeThinkBracketMatch = content.match(/\[THINK\]([\s\S]*?)\[\/THINK\]/);
+
+	if (completeThinkMatch) {
+		const thinkingContent = completeThinkMatch[1]?.trim() ?? '';
+		const cleanContent = `${content.slice(0, completeThinkMatch.index ?? 0)}${content.slice(
+			(completeThinkMatch.index ?? 0) + completeThinkMatch[0].length
+		)}`.trim();
+
+		return {
+			thinking: thinkingContent,
+			cleanContent
+		};
+	}
+
+	if (completeThinkBracketMatch) {
+		const thinkingContent = completeThinkBracketMatch[1]?.trim() ?? '';
+		const cleanContent = `${content.slice(0, completeThinkBracketMatch.index ?? 0)}${content.slice(
+			(completeThinkBracketMatch.index ?? 0) + completeThinkBracketMatch[0].length
+		)}`.trim();
+
+		return {
+			thinking: thinkingContent,
+			cleanContent
 		};
 		};
 	}
 	}
 
 
@@ -39,26 +67,33 @@ export function parseThinkingContent(content: string): {
 }
 }
 
 
 /**
 /**
- * Checks if content contains an opening <think> tag (for streaming)
+ * Checks if content contains an opening thinking tag (for streaming)
+ * Supports both <think> and [THINK] formats
  * @param content - The message content to check
  * @param content - The message content to check
- * @returns True if the content contains an opening <think> tag
+ * @returns True if the content contains an opening thinking tag
  */
  */
 export function hasThinkingStart(content: string): boolean {
 export function hasThinkingStart(content: string): boolean {
-	return content.includes('<think>') || content.includes('<|channel|>analysis');
+	return (
+		content.includes('<think>') ||
+		content.includes('[THINK]') ||
+		content.includes('<|channel|>analysis')
+	);
 }
 }
 
 
 /**
 /**
- * Checks if content contains a closing </think> tag (for streaming)
+ * Checks if content contains a closing thinking tag (for streaming)
+ * Supports both </think> and [/THINK] formats
  * @param content - The message content to check
  * @param content - The message content to check
- * @returns True if the content contains a closing </think> tag
+ * @returns True if the content contains a closing thinking tag
  */
  */
 export function hasThinkingEnd(content: string): boolean {
 export function hasThinkingEnd(content: string): boolean {
-	return content.includes('</think>');
+	return content.includes('</think>') || content.includes('[/THINK]');
 }
 }
 
 
 /**
 /**
  * Extracts partial thinking content during streaming
  * Extracts partial thinking content during streaming
- * Used when we have <think> but not yet </think>
+ * Supports both <think> and [THINK] formats
+ * Used when we have opening tag but not yet closing tag
  * @param content - The message content to extract partial thinking from
  * @param content - The message content to extract partial thinking from
  * @returns An object containing the extracted partial thinking content and the remaining content
  * @returns An object containing the extracted partial thinking content and the remaining content
  */
  */
@@ -66,23 +101,41 @@ export function extractPartialThinking(content: string): {
 	thinking: string | null;
 	thinking: string | null;
 	remainingContent: string;
 	remainingContent: string;
 } {
 } {
-	const startIndex = content.indexOf('<think>');
-	if (startIndex === -1) {
-		return { thinking: null, remainingContent: content };
-	}
+	const thinkStartIndex = content.indexOf('<think>');
+	const thinkEndIndex = content.indexOf('</think>');
 
 
-	const endIndex = content.indexOf('</think>');
-	if (endIndex === -1) {
-		// Still streaming thinking content
-		const thinkingStart = startIndex + '<think>'.length;
-		return {
-			thinking: content.substring(thinkingStart),
-			remainingContent: content.substring(0, startIndex)
-		};
+	const bracketStartIndex = content.indexOf('[THINK]');
+	const bracketEndIndex = content.indexOf('[/THINK]');
+
+	const useThinkFormat =
+		thinkStartIndex !== -1 && (bracketStartIndex === -1 || thinkStartIndex < bracketStartIndex);
+	const useBracketFormat =
+		bracketStartIndex !== -1 && (thinkStartIndex === -1 || bracketStartIndex < thinkStartIndex);
+
+	if (useThinkFormat) {
+		if (thinkEndIndex === -1) {
+			const thinkingStart = thinkStartIndex + '<think>'.length;
+
+			return {
+				thinking: content.substring(thinkingStart),
+				remainingContent: content.substring(0, thinkStartIndex)
+			};
+		}
+	} else if (useBracketFormat) {
+		if (bracketEndIndex === -1) {
+			const thinkingStart = bracketStartIndex + '[THINK]'.length;
+
+			return {
+				thinking: content.substring(thinkingStart),
+				remainingContent: content.substring(0, bracketStartIndex)
+			};
+		}
+	} else {
+		return { thinking: null, remainingContent: content };
 	}
 	}
 
 
-	// Complete thinking block found
 	const parsed = parseThinkingContent(content);
 	const parsed = parseThinkingContent(content);
+
 	return {
 	return {
 		thinking: parsed.thinking,
 		thinking: parsed.thinking,
 		remainingContent: parsed.cleanContent
 		remainingContent: parsed.cleanContent

+ 166 - 0
tools/server/webui/src/stories/ChatMessage.stories.svelte

@@ -59,6 +59,60 @@
 		thinking: '',
 		thinking: '',
 		children: []
 		children: []
 	});
 	});
+
+	// Message with <think> format thinking content
+	const thinkTagMessage: DatabaseMessage = {
+		id: '6',
+		convId: 'conv-1',
+		type: 'message',
+		timestamp: Date.now() - 1000 * 60 * 2,
+		role: 'assistant',
+		content:
+			"<think>\nLet me analyze this step by step:\n\n1. The user is asking about thinking formats\n2. I need to demonstrate the &lt;think&gt; tag format\n3. This content should be displayed in the thinking section\n4. The main response should be separate\n\nThis is a good example of reasoning content.\n</think>\n\nHere's my response after thinking through the problem. The thinking content above should be displayed separately from this main response content.",
+		parent: '1',
+		thinking: '',
+		children: []
+	};
+
+	// Message with [THINK] format thinking content
+	const thinkBracketMessage: DatabaseMessage = {
+		id: '7',
+		convId: 'conv-1',
+		type: 'message',
+		timestamp: Date.now() - 1000 * 60 * 1,
+		role: 'assistant',
+		content:
+			'[THINK]\nThis is the DeepSeek-style thinking format:\n\n- Using square brackets instead of angle brackets\n- Should work identically to the &lt;think&gt; format\n- Content parsing should extract this reasoning\n- Display should be the same as &lt;think&gt; format\n\nBoth formats should be supported seamlessly.\n[/THINK]\n\nThis is the main response content that comes after the [THINK] block. The reasoning above should be parsed and displayed in the thinking section.',
+		parent: '1',
+		thinking: '',
+		children: []
+	};
+
+	// Streaming message for <think> format
+	let streamingThinkMessage = $state({
+		id: '8',
+		convId: 'conv-1',
+		type: 'message',
+		timestamp: 0, // No timestamp = streaming
+		role: 'assistant',
+		content: '',
+		parent: '1',
+		thinking: '',
+		children: []
+	});
+
+	// Streaming message for [THINK] format
+	let streamingBracketMessage = $state({
+		id: '9',
+		convId: 'conv-1',
+		type: 'message',
+		timestamp: 0, // No timestamp = streaming
+		role: 'assistant',
+		content: '',
+		parent: '1',
+		thinking: '',
+		children: []
+	});
 </script>
 </script>
 
 
 <Story
 <Story
@@ -144,3 +198,115 @@
 		await new Promise(resolve => setTimeout(resolve, 100));
 		await new Promise(resolve => setTimeout(resolve, 100));
 	}}
 	}}
 />
 />
+
+<Story
+	name="ThinkTagFormat"
+	args={{
+		class: 'max-w-[56rem] w-[calc(100vw-2rem)]',
+		message: thinkTagMessage
+	}}
+/>
+
+<Story
+	name="ThinkBracketFormat"
+	args={{
+		class: 'max-w-[56rem] w-[calc(100vw-2rem)]',
+		message: thinkBracketMessage
+	}}
+/>
+
+<Story
+	name="StreamingThinkTag"
+	args={{
+		message: streamingThinkMessage
+	}}
+	parameters={{
+		test: {
+			timeout: 30000
+		}
+	}}
+	asChild
+	play={async () => {
+		// Phase 1: Stream <think> reasoning content
+		const thinkingContent =
+			'Let me work through this problem systematically:\n\n1. First, I need to understand what the user is asking\n2. Then I should consider different approaches\n3. I need to evaluate the pros and cons\n4. Finally, I should provide a clear recommendation\n\nThis step-by-step approach will ensure accuracy.';
+
+		let currentContent = '<think>\n';
+		streamingThinkMessage.content = currentContent;
+
+		for (let i = 0; i < thinkingContent.length; i++) {
+			currentContent += thinkingContent[i];
+			streamingThinkMessage.content = currentContent;
+			await new Promise((resolve) => setTimeout(resolve, 5));
+		}
+
+		// Close the thinking block
+		currentContent += '\n</think>\n\n';
+		streamingThinkMessage.content = currentContent;
+		await new Promise((resolve) => setTimeout(resolve, 200));
+
+		// Phase 2: Stream main response content
+		const responseContent =
+			"Based on my analysis above, here's the solution:\n\n**Key Points:**\n- The approach should be systematic\n- We need to consider all factors\n- Implementation should be step-by-step\n\nThis ensures the best possible outcome.";
+
+		for (let i = 0; i < responseContent.length; i++) {
+			currentContent += responseContent[i];
+			streamingThinkMessage.content = currentContent;
+			await new Promise((resolve) => setTimeout(resolve, 10));
+		}
+
+		streamingThinkMessage.timestamp = Date.now();
+	}}
+>
+	<div class="w-[56rem]">
+		<ChatMessage message={streamingThinkMessage} />
+	</div>
+</Story>
+
+<Story
+	name="StreamingThinkBracket"
+	args={{
+		message: streamingBracketMessage
+	}}
+	parameters={{
+		test: {
+			timeout: 30000
+		}
+	}}
+	asChild
+	play={async () => {
+		// Phase 1: Stream [THINK] reasoning content
+		const thinkingContent =
+			'Using the DeepSeek format now:\n\n- This demonstrates the &#91;THINK&#93; bracket format\n- Should parse identically to &lt;think&gt; tags\n- The UI should display this in the thinking section\n- Main content should be separate\n\nBoth formats provide the same functionality.';
+
+		let currentContent = '[THINK]\n';
+		streamingBracketMessage.content = currentContent;
+
+		for (let i = 0; i < thinkingContent.length; i++) {
+			currentContent += thinkingContent[i];
+			streamingBracketMessage.content = currentContent;
+			await new Promise((resolve) => setTimeout(resolve, 5));
+		}
+
+		// Close the thinking block
+		currentContent += '\n[/THINK]\n\n';
+		streamingBracketMessage.content = currentContent;
+		await new Promise((resolve) => setTimeout(resolve, 200));
+
+		// Phase 2: Stream main response content
+		const responseContent =
+			"Here's my response after using the &#91;THINK&#93; format:\n\n**Observations:**\n- Both &lt;think&gt; and &#91;THINK&#93; formats work seamlessly\n- The parsing logic handles both cases\n- UI display is consistent across formats\n\nThis demonstrates the enhanced thinking content support.";
+
+		for (let i = 0; i < responseContent.length; i++) {
+			currentContent += responseContent[i];
+			streamingBracketMessage.content = currentContent;
+			await new Promise((resolve) => setTimeout(resolve, 10));
+		}
+
+		streamingBracketMessage.timestamp = Date.now();
+	}}
+>
+	<div class="w-[56rem]">
+		<ChatMessage message={streamingBracketMessage} />
+	</div>
+</Story>