Skip to content

Commit

Permalink
feat: add support for <think> tags in ChatMessageContent
Browse files Browse the repository at this point in the history
This PR adds support for rendering <think> tags in the ChatMessageContent component. The <think> content is displayed in a collapsible details element, allowing users to view the AI's internal thought process.

The MarkdownFromCody component is also updated to allow the <think> element.
  • Loading branch information
abeatrix committed Jan 28, 2025
1 parent 6e8d30d commit ae89419
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 1 deletion.
31 changes: 30 additions & 1 deletion vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,22 @@ interface ChatMessageContentProps {
className?: string
}

const extractThinkContent = (content: string): { displayContent: string; thinkContent: string } => {
const thinkRegex = /<think>([\s\S]*?)<\/think>/g
const thinkMatches = [...content.matchAll(thinkRegex)]

// Collect all think content
const thinkContent = thinkMatches
.map(match => match[1].trim())
.filter(Boolean)
.join('\n\n')

// Remove think tags from display content
const displayContent = content.replace(thinkRegex, '')

return { displayContent, thinkContent }
}

/**
* A component presenting the content of a chat message.
*/
Expand Down Expand Up @@ -204,10 +220,23 @@ export const ChatMessageContent: React.FunctionComponent<ChatMessageContentProps
smartApplyStates,
])

const { displayContent, thinkContent } = useMemo(
() => extractThinkContent(displayMarkdown),
[displayMarkdown]
)

return (
<div ref={rootRef} data-testid="chat-message-content">
{thinkContent && (
<details className="tw-container tw-my-4 tw-border tw-border-muted-foreground tw-rounded-md tw-p-2">
<summary className="tw-px-2">Thinking...</summary>
<MarkdownFromCody className="tw-my-2 tw-text-muted-foreground tw-p-2">
{thinkContent}
</MarkdownFromCody>
</details>
)}
<MarkdownFromCody className={clsx(styles.content, className)}>
{displayMarkdown}
{displayContent}
</MarkdownFromCody>
</div>
)
Expand Down
1 change: 1 addition & 0 deletions vscode/webviews/components/MarkdownFromCody.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ const ALLOWED_ELEMENTS = [
'h5',
'h6',
'br',
'think',
]

function defaultUrlProcessor(url: string): string {
Expand Down

0 comments on commit ae89419

Please sign in to comment.