Skip to content

Commit 606b880

Browse files
author
Sagid Magomedov
committed
fix: parse Gemma 4 <thought> reasoning tags alongside <think>
Gemma 4 streams reasoning inside <thought>...</thought> instead of <think>...</think>. Without this the content leaks into chat text and the agent triggers a retry on the first turn. - TagMatcher: support multiple tag names - string[], track activeTagName so <think> is never closed by </thought> (and vice-versa). - base-openai-compatible-provider and openai handler: match both tags. - Tests: <thought> parsing, cross-tag isolation, and invariants.
1 parent b761a0a commit 606b880

4 files changed

Lines changed: 114 additions & 19 deletions

File tree

src/api/providers/__tests__/base-openai-compatible-provider.spec.ts

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,75 @@ describe("BaseOpenAiCompatibleProvider", () => {
9595
])
9696
})
9797

98+
it("should handle reasoning tags (<thought>) from stream", async () => {
99+
mockCreate.mockImplementationOnce(() => {
100+
return {
101+
[Symbol.asyncIterator]: () => ({
102+
next: vi
103+
.fn()
104+
.mockResolvedValueOnce({
105+
done: false,
106+
value: { choices: [{ delta: { content: "<thought>Deep thought" } }] },
107+
})
108+
.mockResolvedValueOnce({
109+
done: false,
110+
value: { choices: [{ delta: { content: " here</thought>" } }] },
111+
})
112+
.mockResolvedValueOnce({
113+
done: false,
114+
value: { choices: [{ delta: { content: "Result: 42" } }] },
115+
})
116+
.mockResolvedValueOnce({ done: true }),
117+
}),
118+
}
119+
})
120+
const stream = handler.createMessage("system prompt", [])
121+
const chunks = []
122+
for await (const chunk of stream) {
123+
chunks.push(chunk)
124+
}
125+
expect(chunks).toEqual([
126+
{ type: "reasoning", text: "Deep thought" },
127+
{ type: "reasoning", text: " here" },
128+
{ type: "text", text: "Result: 42" },
129+
])
130+
})
131+
132+
it("should not close <think> tag with </thought> tag", async () => {
133+
mockCreate.mockImplementationOnce(() => {
134+
return {
135+
[Symbol.asyncIterator]: () => ({
136+
next: vi
137+
.fn()
138+
.mockResolvedValueOnce({
139+
done: false,
140+
value: { choices: [{ delta: { content: "<think>Thinking" } }] },
141+
})
142+
.mockResolvedValueOnce({
143+
done: false,
144+
value: { choices: [{ delta: { content: " but closing with wrong tag</thought>" } }] },
145+
})
146+
.mockResolvedValueOnce({
147+
done: false,
148+
value: { choices: [{ delta: { content: " still thinking" } }] },
149+
})
150+
.mockResolvedValueOnce({ done: true }),
151+
}),
152+
}
153+
})
154+
const stream = handler.createMessage("system prompt", [])
155+
const chunks = []
156+
for await (const chunk of stream) {
157+
chunks.push(chunk)
158+
}
159+
// The </thought> tag should be treated as text since it doesn't match the active <think> tag
160+
expect(chunks).toEqual([
161+
{ type: "reasoning", text: "Thinking" },
162+
{ type: "reasoning", text: " but closing with wrong tag</thought>" },
163+
{ type: "reasoning", text: " still thinking" },
164+
])
165+
})
166+
98167
it("should handle complete <think> tag in a single chunk", async () => {
99168
mockCreate.mockImplementationOnce(() => {
100169
return {

src/api/providers/base-openai-compatible-provider.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
118118
const stream = await this.createStream(systemPrompt, messages, metadata)
119119

120120
const matcher = new TagMatcher(
121-
"think",
121+
["think", "thought"],
122122
(chunk) =>
123123
({
124124
type: chunk.matched ? "reasoning" : "text",

src/api/providers/openai.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
184184
}
185185

186186
const matcher = new TagMatcher(
187-
"think",
187+
["think", "thought"],
188188
(chunk) =>
189189
({
190190
type: chunk.matched ? "reasoning" : "text",

src/utils/tag-matcher.ts

Lines changed: 43 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,17 @@ export class TagMatcher<Result = TagMatcherResult> {
1717
state: "TEXT" | "TAG_OPEN" | "TAG_CLOSE" = "TEXT"
1818
depth = 0
1919
pointer = 0
20+
private readonly tagNames: string[]
21+
private activeTagName?: string
22+
private candidates: { name: string; index: number }[] = []
23+
2024
constructor(
21-
readonly tagName: string,
25+
tagName: string | string[],
2226
readonly transform?: (chunks: TagMatcherResult) => Result,
2327
readonly position = 0,
24-
) {}
28+
) {
29+
this.tagNames = Array.isArray(tagName) ? tagName : [tagName]
30+
}
2531
private collect() {
2632
if (!this.cached.length) {
2733
return
@@ -57,38 +63,56 @@ export class TagMatcher<Result = TagMatcherResult> {
5763
if (char === "<" && (this.pointer <= this.position + 1 || this.matched)) {
5864
this.state = "TAG_OPEN"
5965
this.index = 0
66+
this.candidates = this.tagNames.map((name) => ({ name, index: 0 }))
6067
} else {
6168
this.collect()
6269
}
6370
} else if (this.state === "TAG_OPEN") {
64-
if (char === ">" && this.index === this.tagName.length) {
65-
this.state = "TEXT"
66-
if (!this.matched) {
67-
this.cached = []
71+
if (char === ">") {
72+
const matched = this.candidates.find((c) => c.index === c.name.length)
73+
if (matched) {
74+
this.state = "TEXT"
75+
this.activeTagName = matched.name
76+
if (!this.matched) {
77+
this.cached = []
78+
}
79+
this.depth++
80+
this.matched = true
81+
continue
6882
}
69-
this.depth++
70-
this.matched = true
71-
} else if (this.index === 0 && char === "/") {
83+
} else if (this.candidates.every((c) => c.index === 0) && char === "/") {
7284
this.state = "TAG_CLOSE"
73-
} else if (char === " " && (this.index === 0 || this.index === this.tagName.length)) {
85+
this.index = 0
7486
continue
75-
} else if (this.tagName[this.index] === char) {
76-
this.index++
87+
} else if (char === " ") {
88+
const remaining = this.candidates.filter((c) => c.index === 0 || c.index === c.name.length)
89+
if (remaining.length === this.candidates.length) {
90+
continue
91+
}
92+
this.candidates = remaining
7793
} else {
78-
this.state = "TEXT"
79-
this.collect()
94+
this.candidates = this.candidates.filter((c) => c.name[c.index] === char)
95+
for (const c of this.candidates) {
96+
c.index++
97+
}
98+
if (this.candidates.length === 0) {
99+
this.state = "TEXT"
100+
this.collect()
101+
}
80102
}
81103
} else if (this.state === "TAG_CLOSE") {
82-
if (char === ">" && this.index === this.tagName.length) {
104+
const tagName = this.activeTagName || this.tagNames[0]
105+
if (char === ">" && this.index === tagName.length) {
83106
this.state = "TEXT"
84107
this.depth--
85108
this.matched = this.depth > 0
86109
if (!this.matched) {
110+
this.activeTagName = undefined
87111
this.cached = []
88112
}
89-
} else if (char === " " && (this.index === 0 || this.index === this.tagName.length)) {
113+
} else if (char === " " && (this.index === 0 || this.index === tagName.length)) {
90114
continue
91-
} else if (this.tagName[this.index] === char) {
115+
} else if (tagName[this.index] === char) {
92116
this.index++
93117
} else {
94118
this.state = "TEXT"
@@ -102,6 +126,8 @@ export class TagMatcher<Result = TagMatcherResult> {
102126
this._update(chunk)
103127
}
104128
this.collect()
129+
this.candidates = []
130+
this.activeTagName = undefined
105131
return this.pop()
106132
}
107133
update(chunk: string) {

0 commit comments

Comments
 (0)