fix: support both <think> and <thinking> tags for LM Studio GPT-OSS models

- Created MultiTagXmlMatcher utility to handle multiple XML tag names
- Updated LM Studio handler to parse both <think> and <thinking> tags
- Added comprehensive tests for the new functionality
- Fixes #6750
This commit is contained in:
Roo Code 2025-08-06 15:11:40 +00:00
parent 2b647ed9a1
commit 924a793446
4 changed files with 328 additions and 3 deletions

View file

@ -114,6 +114,94 @@ describe("LmStudioHandler", () => {
expect(textChunks[0].text).toBe("Test response")
})
it("should handle <think> tags in responses", async () => {
mockCreate.mockImplementationOnce(async (options) => {
return {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [
{
delta: { content: "Before <think>This is a thought</think> After" },
index: 0,
},
],
usage: null,
}
yield {
choices: [
{
delta: {},
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 15,
total_tokens: 25,
},
}
},
}
})
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
const textChunks = chunks.filter((chunk) => chunk.type === "text")
const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
expect(textChunks).toContainEqual({ type: "text", text: "Before " })
expect(textChunks).toContainEqual({ type: "text", text: " After" })
expect(reasoningChunks).toContainEqual({ type: "reasoning", text: "This is a thought" })
})
it("should handle <thinking> tags in responses (GPT-OSS compatibility)", async () => {
mockCreate.mockImplementationOnce(async (options) => {
return {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [
{
delta: { content: "Before <thinking>This is thinking content</thinking> After" },
index: 0,
},
],
usage: null,
}
yield {
choices: [
{
delta: {},
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 20,
total_tokens: 30,
},
}
},
}
})
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
const textChunks = chunks.filter((chunk) => chunk.type === "text")
const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
expect(textChunks).toContainEqual({ type: "text", text: "Before " })
expect(textChunks).toContainEqual({ type: "text", text: " After" })
expect(reasoningChunks).toContainEqual({ type: "reasoning", text: "This is thinking content" })
})
it("should handle API errors", async () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))

View file

@ -6,7 +6,7 @@ import { type ModelInfo, openAiModelInfoSaneDefaults, LMSTUDIO_DEFAULT_TEMPERATU
import type { ApiHandlerOptions } from "../../shared/api"
import { XmlMatcher } from "../../utils/xml-matcher"
import { MultiTagXmlMatcher } from "../../utils/multi-tag-xml-matcher"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
@ -87,8 +87,9 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
const results = await this.client.chat.completions.create(params)
const matcher = new XmlMatcher(
"think",
// Support both <think> and <thinking> tags for different GPT-OSS models
const matcher = new MultiTagXmlMatcher(
["think", "thinking"],
(chunk) =>
({
type: chunk.matched ? "reasoning" : "text",

View file

@ -0,0 +1,95 @@
import { MultiTagXmlMatcher } from "../multi-tag-xml-matcher"
describe("MultiTagXmlMatcher", () => {
it("should match content with <think> tags", () => {
const matcher = new MultiTagXmlMatcher(["think", "thinking"])
const input = "Before <think>This is thinking content</think> After"
const results = matcher.update(input)
const finalResults = matcher.final()
const allResults = [...results, ...finalResults]
// Check that we have thinking content
const thinkingBlocks = allResults.filter((r) => r.matched)
const textBlocks = allResults.filter((r) => !r.matched)
expect(thinkingBlocks).toContainEqual({ matched: true, data: "This is thinking content" })
expect(textBlocks.some((b) => b.data.includes("Before"))).toBe(true)
expect(textBlocks.some((b) => b.data.includes("After"))).toBe(true)
})
it("should match content with <thinking> tags", () => {
const matcher = new MultiTagXmlMatcher(["think", "thinking"])
const input = "Before <thinking>This is thinking content</thinking> After"
const results = matcher.update(input)
const finalResults = matcher.final()
const allResults = [...results, ...finalResults]
// Check that we have thinking content
const thinkingBlocks = allResults.filter((r) => r.matched)
const textBlocks = allResults.filter((r) => !r.matched)
expect(thinkingBlocks).toContainEqual({ matched: true, data: "This is thinking content" })
expect(textBlocks.some((b) => b.data.includes("Before"))).toBe(true)
expect(textBlocks.some((b) => b.data.includes("After"))).toBe(true)
})
it("should handle mixed tags in the same content", () => {
const matcher = new MultiTagXmlMatcher(["think", "thinking"])
const input = "Start <think>First thought</think> Middle <thinking>Second thought</thinking> End"
const results = matcher.update(input)
const finalResults = matcher.final()
const allResults = [...results, ...finalResults]
// The important thing is that both thinking blocks are captured
const thinkingBlocks = allResults.filter((r) => r.matched)
const textBlocks = allResults.filter((r) => !r.matched)
expect(thinkingBlocks).toContainEqual({ matched: true, data: "First thought" })
expect(thinkingBlocks).toContainEqual({ matched: true, data: "Second thought" })
expect(textBlocks.some((b) => b.data.includes("Start"))).toBe(true)
expect(textBlocks.some((b) => b.data.includes("Middle"))).toBe(true)
expect(textBlocks.some((b) => b.data.includes("End"))).toBe(true)
})
it("should work with custom transform function", () => {
const transform = (chunk: any) => ({
type: chunk.matched ? "reasoning" : "text",
text: chunk.data,
})
const matcher = new MultiTagXmlMatcher(["think", "thinking"], transform)
const input = "Before <thinking>Reasoning here</thinking> After"
const results = matcher.update(input)
const finalResults = matcher.final()
const allResults = [...results, ...finalResults]
// Check that transform is applied
const reasoningBlocks = allResults.filter((r) => r.type === "reasoning")
const textBlocks = allResults.filter((r) => r.type === "text")
expect(reasoningBlocks).toContainEqual({ type: "reasoning", text: "Reasoning here" })
expect(textBlocks.length).toBeGreaterThan(0)
})
it("should handle empty tags", () => {
const matcher = new MultiTagXmlMatcher(["think", "thinking"])
const input = "Before <think></think> Middle <thinking></thinking> After"
const results = matcher.update(input)
const finalResults = matcher.final()
const allResults = [...results, ...finalResults]
// Empty tags should still be matched but with empty content
const emptyBlocks = allResults.filter((r) => r.matched && r.data === "")
expect(emptyBlocks.length).toBeGreaterThan(0)
})
})

View file

@ -0,0 +1,141 @@
import { XmlMatcherResult } from "./xml-matcher"
/**
* A multi-tag XML matcher that can match multiple tag names.
* This is useful for handling different thinking tag formats from various models.
*/
export class MultiTagXmlMatcher<Result = XmlMatcherResult> {
private buffer = ""
private chunks: Result[] = []
private state: "TEXT" | "TAG_OPEN" | "TAG_CLOSE" = "TEXT"
private currentTag = ""
private depth = 0
private matchedTag = ""
private matchedContent = ""
private lastEmittedIndex = 0
constructor(
private tagNames: string[],
private transform?: (chunks: XmlMatcherResult) => Result,
private position = 0,
) {}
private emit(matched: boolean, data: string) {
// Allow empty strings for empty tags
const result: XmlMatcherResult = { matched, data }
if (this.transform) {
this.chunks.push(this.transform(result))
} else {
this.chunks.push(result as Result)
}
}
private processBuffer() {
let i = 0
while (i < this.buffer.length) {
const char = this.buffer[i]
if (this.state === "TEXT") {
if (char === "<") {
// Emit any text before the tag
if (i > this.lastEmittedIndex) {
this.emit(false, this.buffer.substring(this.lastEmittedIndex, i))
}
this.state = "TAG_OPEN"
this.currentTag = ""
this.lastEmittedIndex = i
}
} else if (this.state === "TAG_OPEN") {
if (char === ">") {
// Check if this is a closing tag
const isClosing = this.currentTag.startsWith("/")
const tagName = isClosing ? this.currentTag.substring(1) : this.currentTag
if (this.tagNames.includes(tagName)) {
if (isClosing && this.matchedTag === tagName) {
this.depth--
if (this.depth === 0) {
// Emit the matched content
this.emit(true, this.matchedContent)
this.matchedContent = ""
this.matchedTag = ""
this.lastEmittedIndex = i + 1
}
} else if (!isClosing) {
if (this.depth === 0) {
this.matchedTag = tagName
this.lastEmittedIndex = i + 1
this.matchedContent = "" // Reset matched content
}
this.depth++
}
}
this.state = "TEXT"
} else if (char !== "/" || this.currentTag.length > 0) {
this.currentTag += char
} else {
this.currentTag += char
}
}
// If we're inside a matched tag, collect the content
if (this.depth > 0 && this.state === "TEXT" && i >= this.lastEmittedIndex) {
this.matchedContent += char
}
i++
}
// Emit any remaining text
if (this.state === "TEXT" && this.depth === 0 && this.lastEmittedIndex < this.buffer.length) {
this.emit(false, this.buffer.substring(this.lastEmittedIndex))
this.lastEmittedIndex = this.buffer.length
}
}
update(chunk: string): Result[] {
this.chunks = []
this.buffer += chunk
this.processBuffer()
// Keep unprocessed content in buffer
if (this.lastEmittedIndex > 0 && this.depth === 0) {
this.buffer = this.buffer.substring(this.lastEmittedIndex)
this.lastEmittedIndex = 0
}
const result = this.chunks
this.chunks = []
return result
}
final(chunk?: string): Result[] {
this.chunks = []
if (chunk) {
this.buffer += chunk
}
// Process any remaining buffer
this.processBuffer()
// Emit any remaining content
if (this.buffer.length > this.lastEmittedIndex) {
if (this.depth > 0 && this.matchedContent) {
// Incomplete tag, emit as text
this.emit(false, this.buffer.substring(this.lastEmittedIndex))
} else {
this.emit(false, this.buffer.substring(this.lastEmittedIndex))
}
}
// Reset state
this.buffer = ""
this.lastEmittedIndex = 0
this.depth = 0
this.matchedTag = ""
this.matchedContent = ""
this.state = "TEXT"
return this.chunks
}
}