Skip to content

Commit 94ac98e

Browse files
roomoteroomote
committed
feat: add support for Magistral [THINK]...[/THINK] reasoning tags
- Created new SquareBracketMatcher class to handle square bracket syntax - Updated native-ollama provider to detect Magistral models and use appropriate matcher - Added comprehensive tests for square bracket matching functionality - Maintains backward compatibility with existing XML-style <think> tags Fixes #8522
1 parent 85b0e8a commit 94ac98e

File tree

3 files changed

+355
-8
lines changed

3 files changed

+355
-8
lines changed

‎src/api/providers/native-ollama.ts‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { BaseProvider } from "./base-provider"
66
import type { ApiHandlerOptions } from "../../shared/api"
77
import { getOllamaModels } from "./fetchers/ollama"
88
import { XmlMatcher } from "../../utils/xml-matcher"
9+
import { SquareBracketMatcher } from "../../utils/square-bracket-matcher"
910
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
1011

1112
interface OllamaChatOptions {
@@ -173,20 +174,31 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
173174
const client = this.ensureClient()
174175
const { id: modelId, info: modelInfo } = await this.fetchModel()
175176
const useR1Format = modelId.toLowerCase().includes("deepseek-r1")
177+
const useMagistralFormat = modelId.toLowerCase().includes("magistral")
176178

177179
const ollamaMessages: Message[] = [
178180
{ role: "system", content: systemPrompt },
179181
...convertToOllamaMessages(messages),
180182
]
181183

182-
const matcher = new XmlMatcher(
183-
"think",
184-
(chunk) =>
185-
({
186-
type: chunk.matched ? "reasoning" : "text",
187-
text: chunk.data,
188-
}) as const,
189-
)
184+
// Use square bracket matcher for Magistral models, XML matcher for others
185+
const matcher = useMagistralFormat
186+
? new SquareBracketMatcher(
187+
"THINK",
188+
(chunk) =>
189+
({
190+
type: chunk.matched ? "reasoning" : "text",
191+
text: chunk.data,
192+
}) as const,
193+
)
194+
: new XmlMatcher(
195+
"think",
196+
(chunk) =>
197+
({
198+
type: chunk.matched ? "reasoning" : "text",
199+
text: chunk.data,
200+
}) as const,
201+
)
190202

191203
try {
192204
// Build options object conditionally

‎src/utils/__tests__/square-bracket-matcher.spec.ts‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
import { SquareBracketMatcher } from "../square-bracket-matcher"
2+
3+
describe("SquareBracketMatcher", () => {
4+
it("matches square bracket tags at position 0", () => {
5+
const matcher = new SquareBracketMatcher("THINK")
6+
const chunks = [...matcher.update("[THINK]data[/THINK]"), ...matcher.final()]
7+
expect(chunks).toEqual([
8+
{
9+
data: "data",
10+
matched: true,
11+
},
12+
])
13+
})
14+
15+
it("handles uppercase tag names", () => {
16+
const matcher = new SquareBracketMatcher("THINK")
17+
const chunks = [...matcher.update("[THINK]reasoning content[/THINK]"), ...matcher.final()]
18+
expect(chunks).toEqual([
19+
{
20+
data: "reasoning content",
21+
matched: true,
22+
},
23+
])
24+
})
25+
26+
it("handles lowercase tag names", () => {
27+
const matcher = new SquareBracketMatcher("think")
28+
const chunks = [...matcher.update("[think]reasoning content[/think]"), ...matcher.final()]
29+
expect(chunks).toEqual([
30+
{
31+
data: "reasoning content",
32+
matched: true,
33+
},
34+
])
35+
})
36+
37+
it("handles mixed content", () => {
38+
const matcher = new SquareBracketMatcher("THINK")
39+
const chunks = [...matcher.update("Before [THINK]reasoning[/THINK] After"), ...matcher.final()]
40+
expect(chunks).toEqual([
41+
{
42+
data: "Before ",
43+
matched: false,
44+
},
45+
{
46+
data: "reasoning",
47+
matched: true,
48+
},
49+
{
50+
data: " After",
51+
matched: false,
52+
},
53+
])
54+
})
55+
56+
it("handles streaming push", () => {
57+
const matcher = new SquareBracketMatcher("THINK")
58+
const chunks = [
59+
...matcher.update("["),
60+
...matcher.update("THINK"),
61+
...matcher.update("]"),
62+
...matcher.update("reasoning"),
63+
...matcher.update(" content"),
64+
...matcher.update("[/"),
65+
...matcher.update("THINK"),
66+
...matcher.update("]"),
67+
...matcher.final(),
68+
]
69+
expect(chunks).toEqual([
70+
{
71+
data: "reasoning content",
72+
matched: true,
73+
},
74+
])
75+
})
76+
77+
it("handles nested tags", () => {
78+
const matcher = new SquareBracketMatcher("THINK")
79+
const chunks = [...matcher.update("[THINK]X[THINK]Y[/THINK]Z[/THINK]"), ...matcher.final()]
80+
expect(chunks).toEqual([
81+
{
82+
data: "X[THINK]Y[/THINK]Z",
83+
matched: true,
84+
},
85+
])
86+
})
87+
88+
it("handles invalid tag format", () => {
89+
const matcher = new SquareBracketMatcher("THINK")
90+
const chunks = [...matcher.update("[INVALID]data[/INVALID]"), ...matcher.final()]
91+
expect(chunks).toEqual([
92+
{
93+
data: "[INVALID]data[/INVALID]",
94+
matched: false,
95+
},
96+
])
97+
})
98+
99+
it("handles unclosed tags", () => {
100+
const matcher = new SquareBracketMatcher("THINK")
101+
const chunks = [...matcher.update("[THINK]data"), ...matcher.final()]
102+
expect(chunks).toEqual([
103+
{
104+
data: "data",
105+
matched: true,
106+
},
107+
])
108+
})
109+
110+
it("handles wrong matching position", () => {
111+
const matcher = new SquareBracketMatcher("THINK")
112+
const chunks = [...matcher.update("prefix[THINK]data[/THINK]"), ...matcher.final()]
113+
expect(chunks).toEqual([
114+
{
115+
data: "prefix",
116+
matched: false,
117+
},
118+
{
119+
data: "data",
120+
matched: true,
121+
},
122+
])
123+
})
124+
125+
it("handles multiple sequential tags", () => {
126+
const matcher = new SquareBracketMatcher("THINK")
127+
const chunks = [...matcher.update("[THINK]first[/THINK] middle [THINK]second[/THINK]"), ...matcher.final()]
128+
expect(chunks).toEqual([
129+
{
130+
data: "first",
131+
matched: true,
132+
},
133+
{
134+
data: " middle ",
135+
matched: false,
136+
},
137+
{
138+
data: "second",
139+
matched: true,
140+
},
141+
])
142+
})
143+
144+
it("transforms output when transform function is provided", () => {
145+
const matcher = new SquareBracketMatcher("THINK", (chunk) => ({
146+
type: chunk.matched ? "reasoning" : "text",
147+
text: chunk.data,
148+
}))
149+
const chunks = [...matcher.update("Before [THINK]reasoning[/THINK] After"), ...matcher.final()]
150+
expect(chunks).toEqual([
151+
{
152+
type: "text",
153+
text: "Before ",
154+
},
155+
{
156+
type: "reasoning",
157+
text: "reasoning",
158+
},
159+
{
160+
type: "text",
161+
text: " After",
162+
},
163+
])
164+
})
165+
166+
it("handles Magistral-style reasoning blocks", () => {
167+
const matcher = new SquareBracketMatcher("THINK")
168+
const input = `Let me analyze this problem.
169+
170+
[THINK]
171+
I need to understand what the user is asking for.
172+
They want to fix an issue with Magistral model's reasoning tags.
173+
The tags use square brackets instead of angle brackets.
174+
[/THINK]
175+
176+
Based on my analysis, here's the solution...`
177+
178+
const chunks = [...matcher.update(input), ...matcher.final()]
179+
expect(chunks).toEqual([
180+
{
181+
data: "Let me analyze this problem.\n\n",
182+
matched: false,
183+
},
184+
{
185+
data: "\nI need to understand what the user is asking for.\nThey want to fix an issue with Magistral model's reasoning tags.\nThe tags use square brackets instead of angle brackets.\n",
186+
matched: true,
187+
},
188+
{
189+
data: "\n\nBased on my analysis, here's the solution...",
190+
matched: false,
191+
},
192+
])
193+
})
194+
})

‎src/utils/square-bracket-matcher.ts‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
export interface SquareBracketMatcherResult {
2+
matched: boolean
3+
data: string
4+
}
5+
6+
/**
7+
* Matcher for square bracket tags like [THINK]...[/THINK]
8+
* Used by models like Magistral that use square bracket syntax instead of angle brackets
9+
*/
10+
export class SquareBracketMatcher<Result = SquareBracketMatcherResult> {
11+
private buffer = ""
12+
private insideTag = false
13+
private tagDepth = 0
14+
private results: Result[] = []
15+
16+
constructor(
17+
readonly tagName: string,
18+
readonly transform?: (chunks: SquareBracketMatcherResult) => Result,
19+
readonly position = 0,
20+
) {}
21+
22+
private emit(matched: boolean, data: string) {
23+
if (!data) return
24+
const chunk: SquareBracketMatcherResult = { matched, data }
25+
this.results.push(this.transform ? this.transform(chunk) : (chunk as Result))
26+
}
27+
28+
private processComplete() {
29+
const openTag = `[${this.tagName}]`
30+
const closeTag = `[/${this.tagName}]`
31+
let processed = false
32+
33+
while (true) {
34+
if (!this.insideTag) {
35+
// Look for opening tag
36+
const openIndex = this.buffer.indexOf(openTag)
37+
if (openIndex === -1) {
38+
// No opening tag found
39+
break
40+
}
41+
42+
if (openIndex > 0) {
43+
// Emit text before tag
44+
this.emit(false, this.buffer.substring(0, openIndex))
45+
this.buffer = this.buffer.substring(openIndex)
46+
processed = true
47+
}
48+
49+
// Now we have opening tag at start
50+
this.buffer = this.buffer.substring(openTag.length)
51+
this.insideTag = true
52+
this.tagDepth = 1
53+
processed = true
54+
} else {
55+
// Inside tag, look for closing tag
56+
let pos = 0
57+
let contentStart = 0
58+
59+
while (pos < this.buffer.length) {
60+
const nextOpen = this.buffer.indexOf(openTag, pos)
61+
const nextClose = this.buffer.indexOf(closeTag, pos)
62+
63+
if (nextClose === -1) {
64+
// No closing tag found yet
65+
break
66+
}
67+
68+
if (nextOpen !== -1 && nextOpen < nextClose) {
69+
// Found nested opening tag
70+
this.tagDepth++
71+
pos = nextOpen + openTag.length
72+
} else {
73+
// Found closing tag
74+
this.tagDepth--
75+
if (this.tagDepth === 0) {
76+
// Complete match found
77+
const content = this.buffer.substring(contentStart, nextClose)
78+
this.emit(true, content)
79+
this.buffer = this.buffer.substring(nextClose + closeTag.length)
80+
this.insideTag = false
81+
processed = true
82+
break
83+
} else {
84+
// Still nested
85+
pos = nextClose + closeTag.length
86+
}
87+
}
88+
}
89+
90+
if (this.insideTag) {
91+
// Still inside tag, no complete match yet
92+
break
93+
}
94+
}
95+
}
96+
97+
return processed
98+
}
99+
100+
update(chunk: string): Result[] {
101+
this.buffer += chunk
102+
this.results = []
103+
104+
// Process any complete tag pairs
105+
this.processComplete()
106+
107+
// For streaming, only emit unmatched text if we're sure it won't be part of a tag
108+
if (!this.insideTag && this.buffer && !this.buffer.includes("[")) {
109+
// No potential tags, emit as text
110+
this.emit(false, this.buffer)
111+
this.buffer = ""
112+
}
113+
114+
const results = this.results
115+
this.results = []
116+
return results
117+
}
118+
119+
final(chunk?: string): Result[] {
120+
if (chunk) {
121+
this.buffer += chunk
122+
}
123+
124+
this.results = []
125+
126+
// Process any remaining complete pairs
127+
this.processComplete()
128+
129+
// Emit any remaining buffer
130+
if (this.buffer) {
131+
this.emit(this.insideTag, this.buffer)
132+
this.buffer = ""
133+
}
134+
135+
const results = this.results
136+
this.results = []
137+
this.insideTag = false
138+
this.tagDepth = 0
139+
return results
140+
}
141+
}

0 commit comments

Comments
 (0)