Skip to content

Commit db33696

Browse files
mrubensmrubens
authored
Merge pull request #1080 from System233/feat-ollama-deepseek-reasoning
feat: ollama-deepseek reasoning support
2 parents bfd5024 + 0eecda2 commit db33696

File tree

3 files changed

+244
-3
lines changed

3 files changed

+244
-3
lines changed

‎src/api/providers/ollama.ts‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { convertToOpenAiMessages } from "../transform/openai-format"
66
import { convertToR1Format } from "../transform/r1-format"
77
import { ApiStream } from "../transform/stream"
88
import { DEEP_SEEK_DEFAULT_TEMPERATURE } from "./openai"
9+
import { XmlMatcher } from "../../utils/xml-matcher"
910

1011
const OLLAMA_DEFAULT_TEMPERATURE = 0
1112

@@ -35,15 +36,26 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
3536
temperature: this.options.modelTemperature ?? OLLAMA_DEFAULT_TEMPERATURE,
3637
stream: true,
3738
})
39+
const matcher = new XmlMatcher(
40+
"think",
41+
(chunk) =>
42+
({
43+
type: chunk.matched ? "reasoning" : "text",
44+
text: chunk.data,
45+
}) as const,
46+
)
3847
for await (const chunk of stream) {
3948
const delta = chunk.choices[0]?.delta
49+
4050
if (delta?.content) {
41-
yield {
42-
type: "text",
43-
text: delta.content,
51+
for (const chunk of matcher.update(delta.content)) {
52+
yield chunk
4453
}
4554
}
4655
}
56+
for (const chunk of matcher.final()) {
57+
yield chunk
58+
}
4759
}
4860

4961
getModel(): { id: string; info: ModelInfo } {

‎src/utils/__tests__/xml-matcher.test.ts‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
import { XmlMatcher } from "../xml-matcher"
2+
3+
describe("XmlMatcher", () => {
4+
it("only match at position 0", () => {
5+
const matcher = new XmlMatcher("think")
6+
const chunks = [...matcher.update("<think>data</think>"), ...matcher.final()]
7+
expect(chunks).toHaveLength(1)
8+
expect(chunks).toEqual([
9+
{
10+
matched: true,
11+
data: "data",
12+
},
13+
])
14+
})
15+
it("tag with space", () => {
16+
const matcher = new XmlMatcher("think")
17+
const chunks = [...matcher.update("< think >data</ think >"), ...matcher.final()]
18+
expect(chunks).toHaveLength(1)
19+
expect(chunks).toEqual([
20+
{
21+
matched: true,
22+
data: "data",
23+
},
24+
])
25+
})
26+
27+
it("invalid tag", () => {
28+
const matcher = new XmlMatcher("think")
29+
const chunks = [...matcher.update("< think 1>data</ think >"), ...matcher.final()]
30+
expect(chunks).toHaveLength(1)
31+
expect(chunks).toEqual([
32+
{
33+
matched: false,
34+
data: "< think 1>data</ think >",
35+
},
36+
])
37+
})
38+
39+
it("anonymous tag", () => {
40+
const matcher = new XmlMatcher("think")
41+
const chunks = [...matcher.update("<>data</>"), ...matcher.final()]
42+
expect(chunks).toHaveLength(1)
43+
expect(chunks).toEqual([
44+
{
45+
matched: false,
46+
data: "<>data</>",
47+
},
48+
])
49+
})
50+
51+
it("streaming push", () => {
52+
const matcher = new XmlMatcher("think")
53+
const chunks = [
54+
...matcher.update("<thi"),
55+
...matcher.update("nk"),
56+
...matcher.update(">dat"),
57+
...matcher.update("a</"),
58+
...matcher.update("think>"),
59+
]
60+
expect(chunks).toHaveLength(2)
61+
expect(chunks).toEqual([
62+
{
63+
matched: true,
64+
data: "dat",
65+
},
66+
{
67+
matched: true,
68+
data: "a",
69+
},
70+
])
71+
})
72+
73+
it("nested tag", () => {
74+
const matcher = new XmlMatcher("think")
75+
const chunks = [...matcher.update("<think>X<think>Y</think>Z</think>"), ...matcher.final()]
76+
expect(chunks).toHaveLength(1)
77+
expect(chunks).toEqual([
78+
{
79+
matched: true,
80+
data: "X<think>Y</think>Z",
81+
},
82+
])
83+
})
84+
85+
it("nested invalid tag", () => {
86+
const matcher = new XmlMatcher("think")
87+
const chunks = [...matcher.update("<think>X<think>Y</thxink>Z</think>"), ...matcher.final()]
88+
expect(chunks).toHaveLength(2)
89+
expect(chunks).toEqual([
90+
{
91+
matched: true,
92+
data: "X<think>Y</thxink>Z",
93+
},
94+
{
95+
matched: true,
96+
data: "</think>",
97+
},
98+
])
99+
})
100+
101+
it("Wrong matching position", () => {
102+
const matcher = new XmlMatcher("think")
103+
const chunks = [...matcher.update("1<think>data</think>"), ...matcher.final()]
104+
expect(chunks).toHaveLength(1)
105+
expect(chunks).toEqual([
106+
{
107+
matched: false,
108+
data: "1<think>data</think>",
109+
},
110+
])
111+
})
112+
113+
it("Unclosed tag", () => {
114+
const matcher = new XmlMatcher("think")
115+
const chunks = [...matcher.update("<think>data"), ...matcher.final()]
116+
expect(chunks).toHaveLength(1)
117+
expect(chunks).toEqual([
118+
{
119+
matched: true,
120+
data: "data",
121+
},
122+
])
123+
})
124+
})

‎src/utils/xml-matcher.ts‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
export interface XmlMatcherResult {
2+
matched: boolean
3+
data: string
4+
}
5+
export class XmlMatcher<Result = XmlMatcherResult> {
6+
index = 0
7+
chunks: XmlMatcherResult[] = []
8+
cached: string[] = []
9+
matched: boolean = false
10+
state: "TEXT" | "TAG_OPEN" | "TAG_CLOSE" = "TEXT"
11+
depth = 0
12+
pointer = 0
13+
constructor(
14+
readonly tagName: string,
15+
readonly transform?: (chunks: XmlMatcherResult) => Result,
16+
readonly position = 0,
17+
) {}
18+
private collect() {
19+
if (!this.cached.length) {
20+
return
21+
}
22+
const last = this.chunks.at(-1)
23+
const data = this.cached.join("")
24+
const matched = this.matched
25+
if (last?.matched === matched) {
26+
last.data += data
27+
} else {
28+
this.chunks.push({
29+
data,
30+
matched,
31+
})
32+
}
33+
this.cached = []
34+
}
35+
private pop() {
36+
const chunks = this.chunks
37+
this.chunks = []
38+
if (!this.transform) {
39+
return chunks as Result[]
40+
}
41+
return chunks.map(this.transform)
42+
}
43+
44+
private _update(chunk: string) {
45+
for (let i = 0; i < chunk.length; i++) {
46+
const char = chunk[i]
47+
this.cached.push(char)
48+
this.pointer++
49+
50+
if (this.state === "TEXT") {
51+
if (char === "<" && (this.pointer <= this.position + 1 || this.matched)) {
52+
this.state = "TAG_OPEN"
53+
this.index = 0
54+
} else {
55+
this.collect()
56+
}
57+
} else if (this.state === "TAG_OPEN") {
58+
if (char === ">" && this.index === this.tagName.length) {
59+
this.state = "TEXT"
60+
if (!this.matched) {
61+
this.cached = []
62+
}
63+
this.depth++
64+
this.matched = true
65+
} else if (this.index === 0 && char === "/") {
66+
this.state = "TAG_CLOSE"
67+
} else if (char === " " && (this.index === 0 || this.index === this.tagName.length)) {
68+
continue
69+
} else if (this.tagName[this.index] === char) {
70+
this.index++
71+
} else {
72+
this.state = "TEXT"
73+
this.collect()
74+
}
75+
} else if (this.state === "TAG_CLOSE") {
76+
if (char === ">" && this.index === this.tagName.length) {
77+
this.state = "TEXT"
78+
this.depth--
79+
this.matched = this.depth > 0
80+
if (!this.matched) {
81+
this.cached = []
82+
}
83+
} else if (char === " " && (this.index === 0 || this.index === this.tagName.length)) {
84+
continue
85+
} else if (this.tagName[this.index] === char) {
86+
this.index++
87+
} else {
88+
this.state = "TEXT"
89+
this.collect()
90+
}
91+
}
92+
}
93+
}
94+
final(chunk?: string) {
95+
if (chunk) {
96+
this._update(chunk)
97+
}
98+
this.collect()
99+
return this.pop()
100+
}
101+
update(chunk: string) {
102+
this._update(chunk)
103+
return this.pop()
104+
}
105+
}

0 commit comments

Comments
 (0)