From b0ff08a3cfd1402e86ec8ded3d5c51ada13a923f Mon Sep 17 00:00:00 2001
From: duanfuxiang <duanfuxiang@haizhi.com>
Date: Thu, 27 Feb 2025 12:48:05 +0800
Subject: [PATCH] handle <think> tag

---
 src/components/chat-view/ReactMarkdown.tsx |   6 ++
 src/utils/parse-infio-block.test.ts        | 100 +++++++++++++++++++++
 src/utils/parse-infio-block.ts             |  34 +++++++
 3 files changed, 140 insertions(+)

diff --git a/src/components/chat-view/ReactMarkdown.tsx b/src/components/chat-view/ReactMarkdown.tsx
index 8396660..21a8c1b 100644
--- a/src/components/chat-view/ReactMarkdown.tsx
+++ b/src/components/chat-view/ReactMarkdown.tsx
@@ -9,6 +9,7 @@ import {
 
 import MarkdownActionBlock from './MarkdownActionBlock'
 import MarkdownReferenceBlock from './MarkdownReferenceBlock'
+import MarkdownReasoningBlock from './MarkdownReasoningBlock'
 
 function ReactMarkdown({
 	onApply,
@@ -36,6 +37,11 @@ function ReactMarkdown({
 					<Markdown key={index} className="infio-markdown">
 						{block.content}
 					</Markdown>
+				) : block.type === 'think' ? (
+					<MarkdownReasoningBlock 
+						key={index}
+						reasoningContent={block.content} 
+					/>
 				) : block.startLine && block.endLine && block.filename && block.action === InfioBlockAction.Reference ? (
 					<MarkdownReferenceBlock
 						key={index}
diff --git a/src/utils/parse-infio-block.test.ts b/src/utils/parse-infio-block.test.ts
index 7348de1..fec1dc6 100644
--- a/src/utils/parse-infio-block.test.ts
+++ b/src/utils/parse-infio-block.test.ts
@@ -199,4 +199,104 @@ Some text after without closing tag`,
 		const result = parseinfioBlocks(input)
 		expect(result).toEqual(expected)
 	})
+
+	it('should parse a string with think elements', () => {
+		const input = `Some text before
+<think>
+This is a thought that should be parsed separately.
+It might contain multiple lines of text.
+</think>
+Some text after`
+
+		const expected: ParsedInfioBlock[] = [
+			{ type: 'string', content: 'Some text before\n' },
+			{
+				type: 'think',
+				content: `
+This is a thought that should be parsed separately.
+It might contain multiple lines of text.
+`
+			},
+			{ type: 'string', content: '\nSome text after' },
+		]
+
+		const result = parseinfioBlocks(input)
+		expect(result).toEqual(expected)
+	})
+
+	it('should handle empty think elements', () => {
+		const input = `
+      <think></think>
+    `
+
+		const expected: ParsedInfioBlock[] = [
+			{ type: 'string', content: '\n      ' },
+			{
+				type: 'think',
+				content: '',
+			},
+			{ type: 'string', content: '\n    ' },
+		]
+
+		const result = parseinfioBlocks(input)
+		expect(result).toEqual(expected)
+	})
+
+	it('should handle mixed infio_block and think elements', () => {
+		const input = `Start
+<infio_block language="python" filename="script.py">
+def greet(name):
+    print(f"Hello, {name}!")
+</infio_block>
+Middle
+<think>
+Let me think about this problem...
+I need to consider several approaches.
+</think>
+End`
+
+		const expected: ParsedInfioBlock[] = [
+			{ type: 'string', content: 'Start\n' },
+			{
+				type: 'infio_block',
+				content: `
+def greet(name):
+    print(f"Hello, {name}!")
+`,
+				language: 'python',
+				filename: 'script.py',
+			},
+			{ type: 'string', content: '\nMiddle\n' },
+			{
+				type: 'think',
+				content: `
+Let me think about this problem...
+I need to consider several approaches.
+`
+			},
+			{ type: 'string', content: '\nEnd' },
+		]
+
+		const result = parseinfioBlocks(input)
+		expect(result).toEqual(expected)
+	})
+
+	it('should handle unfinished think with only opening tag', () => {
+		const input = `Start
+<think>
+Some unfinished thought
+without closing tag`
+		const expected: ParsedInfioBlock[] = [
+			{ type: 'string', content: 'Start\n' },
+			{
+				type: 'think',
+				content: `
+Some unfinished thought
+without closing tag`,
+			},
+		]
+
+		const result = parseinfioBlocks(input)
+		expect(result).toEqual(expected)
+	})
 })
diff --git a/src/utils/parse-infio-block.ts b/src/utils/parse-infio-block.ts
index 202621c..09e3ba6 100644
--- a/src/utils/parse-infio-block.ts
+++ b/src/utils/parse-infio-block.ts
@@ -17,6 +17,7 @@ export type ParsedInfioBlock =
 		endLine?: number
 		action?: InfioBlockAction
 	}
+	| { type: 'think'; content: string }
 
 function isInfioBlockAction(value: string): value is InfioBlockAction {
   return Object.values<string>(InfioBlockAction).includes(value)
@@ -82,6 +83,39 @@ export function parseinfioBlocks(input: string): ParsedInfioBlock[] {
 				})
 			}
 			lastEndOffset = endOffset
+		} else if (node.nodeName === 'think') {
+			if (!node.sourceCodeLocation) {
+				throw new Error('sourceCodeLocation is undefined')
+			}
+			const startOffset = node.sourceCodeLocation.startOffset
+			const endOffset = node.sourceCodeLocation.endOffset
+			if (startOffset > lastEndOffset) {
+				parsedResult.push({
+					type: 'string',
+					content: input.slice(lastEndOffset, startOffset),
+				})
+			}
+
+			const children = node.childNodes
+			if (children.length === 0) {
+				parsedResult.push({
+					type: 'think',
+					content: '',
+				})
+			} else {
+				const innerContentStartOffset =
+					children[0].sourceCodeLocation?.startOffset
+				const innerContentEndOffset =
+					children[children.length - 1].sourceCodeLocation?.endOffset
+				if (!innerContentStartOffset || !innerContentEndOffset) {
+					throw new Error('sourceCodeLocation is undefined')
+				}
+				parsedResult.push({
+					type: 'think',
+					content: input.slice(innerContentStartOffset, innerContentEndOffset),
+				})
+			}
+			lastEndOffset = endOffset
 		}
 	}
 	if (lastEndOffset < input.length) {