From 924a79344661b72534b1006eb75c456606ca58ae Mon Sep 17 00:00:00 2001 From: Roo Code Date: Wed, 6 Aug 2025 15:11:40 +0000 Subject: [PATCH] fix: support both and tags for LM Studio GPT-OSS models - Created MultiTagXmlMatcher utility to handle multiple XML tag names - Updated LM Studio handler to parse both and tags - Added comprehensive tests for the new functionality - Fixes #6750 --- src/api/providers/__tests__/lmstudio.spec.ts | 88 +++++++++++ src/api/providers/lm-studio.ts | 7 +- .../__tests__/multi-tag-xml-matcher.spec.ts | 95 ++++++++++++ src/utils/multi-tag-xml-matcher.ts | 141 ++++++++++++++++++ 4 files changed, 328 insertions(+), 3 deletions(-) create mode 100644 src/utils/__tests__/multi-tag-xml-matcher.spec.ts create mode 100644 src/utils/multi-tag-xml-matcher.ts diff --git a/src/api/providers/__tests__/lmstudio.spec.ts b/src/api/providers/__tests__/lmstudio.spec.ts index 0adebdeea7..dddbcf2f0e 100644 --- a/src/api/providers/__tests__/lmstudio.spec.ts +++ b/src/api/providers/__tests__/lmstudio.spec.ts @@ -114,6 +114,94 @@ describe("LmStudioHandler", () => { expect(textChunks[0].text).toBe("Test response") }) + it("should handle tags in responses", async () => { + mockCreate.mockImplementationOnce(async (options) => { + return { + [Symbol.asyncIterator]: async function* () { + yield { + choices: [ + { + delta: { content: "Before This is a thought 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 tags in responses (GPT-OSS compatibility)", async () => { + mockCreate.mockImplementationOnce(async (options) => { + return { + [Symbol.asyncIterator]: async function* () { + yield { + choices: [ + { + delta: { content: "Before This is thinking content 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")) diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index 6c49920bd1..b9beae4ab2 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -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 and tags for different GPT-OSS models + const matcher = new MultiTagXmlMatcher( + ["think", "thinking"], (chunk) => ({ type: chunk.matched ? "reasoning" : "text", diff --git a/src/utils/__tests__/multi-tag-xml-matcher.spec.ts b/src/utils/__tests__/multi-tag-xml-matcher.spec.ts new file mode 100644 index 0000000000..0f662ebf5c --- /dev/null +++ b/src/utils/__tests__/multi-tag-xml-matcher.spec.ts @@ -0,0 +1,95 @@ +import { MultiTagXmlMatcher } from "../multi-tag-xml-matcher" + +describe("MultiTagXmlMatcher", () => { + it("should match content with tags", () => { + const matcher = new MultiTagXmlMatcher(["think", "thinking"]) + const input = "Before This is thinking content 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 tags", () => { + const matcher = new MultiTagXmlMatcher(["think", "thinking"]) + const input = "Before This is thinking content 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 First thought Middle Second thought 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 Reasoning here 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 Middle 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) + }) +}) diff --git a/src/utils/multi-tag-xml-matcher.ts b/src/utils/multi-tag-xml-matcher.ts new file mode 100644 index 0000000000..984708c01c --- /dev/null +++ b/src/utils/multi-tag-xml-matcher.ts @@ -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 { + 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 + } +}