diff --git a/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx b/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx index 11b6244b03..43b17885f9 100644 --- a/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx +++ b/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx @@ -253,6 +253,24 @@ describe('MessageRow', () => { expect(screen.getByText('The capital of France is Paris.')).toBeInTheDocument() }) + it('moves Gemma channel thinking into the trace before rendering the final response', () => { + render( + + ) + + expect(screen.getByText('Thinking trace').closest('[data-thinking-state="complete"]')).toHaveTextContent( + 'Check whether the prompt is a test.' + ) + expect(screen.getByText('Test received!').closest('.select-text')).toBeInTheDocument() + expect(screen.queryByText(/<\|channel\|>thought/)).not.toBeInTheDocument() + expect(screen.queryByText(//)).not.toBeInTheDocument() + }) + it('formats final assistant response text as markdown', () => { render( { ]) }) + it('splits Gemma channel thinking from final response text', () => { + expect( + splitAssistantThinking('<|channel|>thoughtCheck whether the prompt is a test.Test received!') + ).toEqual([ + { kind: 'thinking', text: 'Check whether the prompt is a test.', open: false }, + { kind: 'response', text: 'Test received!' } + ]) + }) + + it('splits Gemma channel thinking when the thought marker is missing the leading pipe', () => { + expect(splitAssistantThinking('thoughtCheck facts.Final answer.')).toEqual([ + { kind: 'thinking', text: 'Check facts.', open: false }, + { kind: 'response', text: 'Final answer.' } + ]) + }) + it('keeps an unclosed think segment open for live streams', () => { expect(splitAssistantThinking('Checking facts')).toEqual([ { kind: 'thinking', text: 'Checking facts', open: true } diff --git a/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts b/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts index dfee7697d2..032351d784 100644 --- a/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts +++ b/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts @@ -13,8 +13,15 @@ type SplitAssistantThinkingOptions = { streaming?: boolean } +type TagMatch = { + index: number + tag: string +} + const THINK_OPEN_TAG = '' const THINK_CLOSE_TAG = '' +const GEMMA_THOUGHT_CHANNEL_TAGS = ['<|channel|>thought', 'thought'] +const GEMMA_CHANNEL_BOUNDARY_TAGS = ['<|channel|>', ''] function indexOfTag(value: string, tag: string, fromIndex: number) { for (let index = fromIndex; index <= value.length - tag.length; index += 1) { @@ -24,12 +31,67 @@ function indexOfTag(value: string, tag: string, fromIndex: number) { return -1 } +function findFirstTag(value: string, tags: string[], fromIndex: number): TagMatch | null { + let bestMatch: TagMatch | null = null + + for (const tag of tags) { + const index = indexOfTag(value, tag, fromIndex) + if (index === -1) continue + if (bestMatch === null || index < bestMatch.index) { + bestMatch = { index, tag } + } + } + + return bestMatch +} + +function splitGemmaChannelThinking(body: string): AssistantContentSegment[] | null { + if (findFirstTag(body, GEMMA_THOUGHT_CHANNEL_TAGS, 0) === null) return null + + const segments: AssistantContentSegment[] = [] + let cursor = 0 + + while (cursor < body.length) { + const open = findFirstTag(body, GEMMA_THOUGHT_CHANNEL_TAGS, cursor) + if (open === null) { + const responseText = body.slice(cursor) + if (responseText.length > 0) { + segments.push({ kind: 'response', text: responseText }) + } + break + } + + const responseText = body.slice(cursor, open.index) + if (responseText.length > 0) { + segments.push({ kind: 'response', text: responseText }) + } + + const thinkingStart = open.index + open.tag.length + const nextChannel = findFirstTag(body, GEMMA_CHANNEL_BOUNDARY_TAGS, thinkingStart) + if (nextChannel === null) { + segments.push({ kind: 'thinking', text: body.slice(thinkingStart), open: true }) + break + } + + const thinkingText = body.slice(thinkingStart, nextChannel.index) + if (thinkingText.length > 0) { + segments.push({ kind: 'thinking', text: thinkingText, open: false }) + } + cursor = nextChannel.index + nextChannel.tag.length + } + + return segments +} + export function splitAssistantThinking( body: string, { streaming = false }: SplitAssistantThinkingOptions = {} ): AssistantContentSegment[] { if (body.length === 0) return [] + const gemmaSegments = splitGemmaChannelThinking(body) + if (gemmaSegments !== null) return gemmaSegments + const segments: AssistantContentSegment[] = [] let cursor = 0 let firstSegment = true