mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
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:
parent
2b647ed9a1
commit
924a793446
4 changed files with 328 additions and 3 deletions
|
|
@ -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"))
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
95
src/utils/__tests__/multi-tag-xml-matcher.spec.ts
Normal file
95
src/utils/__tests__/multi-tag-xml-matcher.spec.ts
Normal 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)
|
||||
})
|
||||
})
|
||||
141
src/utils/multi-tag-xml-matcher.ts
Normal file
141
src/utils/multi-tag-xml-matcher.ts
Normal 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
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue