From bb86eb865142dc1782955890ccc2c47249b76f24 Mon Sep 17 00:00:00 2001 From: Daniel <57051444+daniel-lxs@users.noreply.github.com> Date: Mon, 9 Feb 2026 16:47:49 -0500 Subject: [PATCH] refactor: migrate LM Studio provider to Vercel AI SDK (#11347) * refactor: migrate LM Studio provider to Vercel AI SDK Migrate LmStudioHandler from raw OpenAI SDK to Vercel AI SDK via OpenAICompatibleHandler base class. Changes: - Extend OpenAICompatibleHandler instead of BaseProvider - Use createOpenAICompatible from @ai-sdk/openai-compatible - Use streamText/generateText from ai package - Add extractReasoningMiddleware for tag extraction parity - Pass draft_model via providerOptions for speculative decoding - Remove unused getLmStudioModels function (active version in fetchers/) - Update all tests to mock AI SDK instead of OpenAI SDK * fix: wrap completePrompt with handleAiSdkError for consistent error handling --------- Co-authored-by: Roo Code --- .../__tests__/lm-studio-timeout.spec.ts | 100 ++--- .../__tests__/lmstudio-native-tools.spec.ts | 423 +++++++----------- src/api/providers/__tests__/lmstudio.spec.ts | 289 +++++++++--- src/api/providers/lm-studio.ts | 270 ++++------- 4 files changed, 498 insertions(+), 584 deletions(-) diff --git a/src/api/providers/__tests__/lm-studio-timeout.spec.ts b/src/api/providers/__tests__/lm-studio-timeout.spec.ts index 659fcaaf67..30cf210783 100644 --- a/src/api/providers/__tests__/lm-studio-timeout.spec.ts +++ b/src/api/providers/__tests__/lm-studio-timeout.spec.ts @@ -1,62 +1,69 @@ // npx vitest run api/providers/__tests__/lm-studio-timeout.spec.ts +const { mockCreateOpenAICompatible } = vi.hoisted(() => ({ + mockCreateOpenAICompatible: vi.fn(() => { + return vi.fn(() => ({ + modelId: "llama2", + provider: "lmstudio", + })) + }), +})) + +vi.mock("@ai-sdk/openai-compatible", () => ({ + createOpenAICompatible: mockCreateOpenAICompatible, +})) + +vi.mock("ai", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + streamText: vi.fn(), + generateText: vi.fn(), + } +}) + import { LmStudioHandler } from "../lm-studio" import { ApiHandlerOptions } from "../../../shared/api" -// Mock the timeout config utility -vitest.mock("../utils/timeout-config", () => ({ - getApiRequestTimeout: vitest.fn(), -})) - -import { getApiRequestTimeout } from "../utils/timeout-config" - -// Mock OpenAI -const mockOpenAIConstructor = vitest.fn() -vitest.mock("openai", () => { - return { - __esModule: true, - default: vitest.fn().mockImplementation((config) => { - mockOpenAIConstructor(config) - return { - chat: { - completions: { - create: vitest.fn(), - }, - }, - } - }), - } -}) - -describe("LmStudioHandler timeout configuration", () => { +describe("LmStudioHandler configuration", () => { beforeEach(() => { - vitest.clearAllMocks() + vi.clearAllMocks() }) - it("should use default timeout of 600 seconds when no configuration is set", () => { - ;(getApiRequestTimeout as any).mockReturnValue(600000) - + it("should configure the provider with default base URL", () => { const options: ApiHandlerOptions = { apiModelId: "llama2", lmStudioModelId: "llama2", - lmStudioBaseUrl: "http://localhost:1234", } new LmStudioHandler(options) - expect(getApiRequestTimeout).toHaveBeenCalled() - expect(mockOpenAIConstructor).toHaveBeenCalledWith( + expect(mockCreateOpenAICompatible).toHaveBeenCalledWith( expect.objectContaining({ + name: "lmstudio", baseURL: "http://localhost:1234/v1", apiKey: "noop", - timeout: 600000, // 600 seconds in milliseconds }), ) }) - it("should use custom timeout when configuration is set", () => { - ;(getApiRequestTimeout as any).mockReturnValue(1200000) // 20 minutes + it("should configure the provider with custom base URL", () => { + const options: ApiHandlerOptions = { + apiModelId: "llama2", + lmStudioModelId: "llama2", + lmStudioBaseUrl: "http://localhost:5678", + } + new LmStudioHandler(options) + + expect(mockCreateOpenAICompatible).toHaveBeenCalledWith( + expect.objectContaining({ + baseURL: "http://localhost:5678/v1", + }), + ) + }) + + it("should use 'noop' as the API key", () => { const options: ApiHandlerOptions = { apiModelId: "llama2", lmStudioModelId: "llama2", @@ -65,26 +72,9 @@ describe("LmStudioHandler timeout configuration", () => { new LmStudioHandler(options) - expect(mockOpenAIConstructor).toHaveBeenCalledWith( + expect(mockCreateOpenAICompatible).toHaveBeenCalledWith( expect.objectContaining({ - timeout: 1200000, // 1200 seconds in milliseconds - }), - ) - }) - - it("should handle zero timeout (no timeout)", () => { - ;(getApiRequestTimeout as any).mockReturnValue(0) - - const options: ApiHandlerOptions = { - apiModelId: "llama2", - lmStudioModelId: "llama2", - } - - new LmStudioHandler(options) - - expect(mockOpenAIConstructor).toHaveBeenCalledWith( - expect.objectContaining({ - timeout: 0, // No timeout + apiKey: "noop", }), ) }) diff --git a/src/api/providers/__tests__/lmstudio-native-tools.spec.ts b/src/api/providers/__tests__/lmstudio-native-tools.spec.ts index cca543a269..f9a0b10fb6 100644 --- a/src/api/providers/__tests__/lmstudio-native-tools.spec.ts +++ b/src/api/providers/__tests__/lmstudio-native-tools.spec.ts @@ -1,22 +1,28 @@ // npx vitest run api/providers/__tests__/lmstudio-native-tools.spec.ts -// Mock OpenAI client - must come before other imports -const mockCreate = vi.fn() -vi.mock("openai", () => { +// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls +const { mockStreamText } = vi.hoisted(() => ({ + mockStreamText: vi.fn(), +})) + +vi.mock("ai", async (importOriginal) => { + const actual = await importOriginal() return { - __esModule: true, - default: vi.fn().mockImplementation(() => ({ - chat: { - completions: { - create: mockCreate, - }, - }, - })), + ...actual, + streamText: mockStreamText, } }) +vi.mock("@ai-sdk/openai-compatible", () => ({ + createOpenAICompatible: vi.fn(() => { + return vi.fn(() => ({ + modelId: "local-model", + provider: "lmstudio", + })) + }), +})) + import { LmStudioHandler } from "../lm-studio" -import { NativeToolCallParser } from "../../../core/assistant-message/NativeToolCallParser" import type { ApiHandlerOptions } from "../../../shared/api" describe("LmStudioHandler Native Tools", () => { @@ -49,128 +55,76 @@ describe("LmStudioHandler Native Tools", () => { lmStudioBaseUrl: "http://localhost:1234", } handler = new LmStudioHandler(mockOptions) - - // Clear NativeToolCallParser state before each test - NativeToolCallParser.clearRawChunkState() }) describe("Native Tool Calling Support", () => { - it("should include tools in request when model supports native tools and tools are provided", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + it("should include tools in request when tools are provided", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test response" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }), + }) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", tools: testTools, }) - await stream.next() + for await (const _chunk of stream) { + // consume stream + } - expect(mockCreate).toHaveBeenCalledWith( - expect.objectContaining({ - tools: expect.arrayContaining([ - expect.objectContaining({ - type: "function", - function: expect.objectContaining({ - name: "test_tool", - }), - }), - ]), - }), - ) - // parallel_tool_calls should be true by default when not explicitly set - const callArgs = mockCreate.mock.calls[0][0] - expect(callArgs).toHaveProperty("parallel_tool_calls", true) + const callArgs = mockStreamText.mock.calls[0][0] + expect(callArgs.tools).toBeDefined() }) - it("should include tool_choice when provided", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + it("should include toolChoice when provided", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test response" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }), + }) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", tools: testTools, tool_choice: "auto", }) - await stream.next() + for await (const _chunk of stream) { + // consume stream + } - expect(mockCreate).toHaveBeenCalledWith( - expect.objectContaining({ - tool_choice: "auto", - }), - ) + const callArgs = mockStreamText.mock.calls[0][0] + expect(callArgs.toolChoice).toBe("auto") }) - it("should always include tools and tool_choice in request (tools are always present after PR #10841)", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + it("should yield tool_call_start, tool_call_delta, and tool_call_end chunks from AI SDK stream", async () => { + async function* mockFullStream() { + yield { + type: "tool-input-start", + id: "call_lmstudio_123", + toolName: "test_tool", + } + yield { + type: "tool-input-delta", + id: "call_lmstudio_123", + delta: '{"arg1":"value"}', + } + yield { + type: "tool-input-end", + id: "call_lmstudio_123", + } + } - const stream = handler.createMessage("test prompt", [], { - taskId: "test-task-id", + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }), }) - await stream.next() - - const callArgs = mockCreate.mock.calls[mockCreate.mock.calls.length - 1][0] - // Tools are now always present (minimum 6 from ALWAYS_AVAILABLE_TOOLS) - expect(callArgs).toHaveProperty("tools") - expect(callArgs).toHaveProperty("tool_choice") - // parallel_tool_calls should be true by default when not explicitly set - expect(callArgs).toHaveProperty("parallel_tool_calls", true) - }) - - it("should yield tool_call_partial chunks during streaming", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [ - { - delta: { - tool_calls: [ - { - index: 0, - id: "call_lmstudio_123", - function: { - name: "test_tool", - arguments: '{"arg1":', - }, - }, - ], - }, - }, - ], - } - yield { - choices: [ - { - delta: { - tool_calls: [ - { - index: 0, - function: { - arguments: '"value"}', - }, - }, - ], - }, - }, - ], - } - }, - })) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -182,168 +136,56 @@ describe("LmStudioHandler Native Tools", () => { chunks.push(chunk) } - expect(chunks).toContainEqual({ - type: "tool_call_partial", - index: 0, + const startChunks = chunks.filter((chunk) => chunk.type === "tool_call_start") + expect(startChunks).toHaveLength(1) + expect(startChunks[0]).toEqual({ + type: "tool_call_start", id: "call_lmstudio_123", name: "test_tool", - arguments: '{"arg1":', }) - expect(chunks).toContainEqual({ - type: "tool_call_partial", - index: 0, - id: undefined, - name: undefined, - arguments: '"value"}', - }) - }) - - it("should set parallel_tool_calls based on metadata", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) - - const stream = handler.createMessage("test prompt", [], { - taskId: "test-task-id", - tools: testTools, - parallelToolCalls: true, - }) - await stream.next() - - expect(mockCreate).toHaveBeenCalledWith( - expect.objectContaining({ - parallel_tool_calls: true, - }), - ) - }) - - it("should yield tool_call_end events when finish_reason is tool_calls", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [ - { - delta: { - tool_calls: [ - { - index: 0, - id: "call_lmstudio_test", - function: { - name: "test_tool", - arguments: '{"arg1":"value"}', - }, - }, - ], - }, - }, - ], - } - yield { - choices: [ - { - delta: {}, - finish_reason: "tool_calls", - }, - ], - } - }, - })) - - const stream = handler.createMessage("test prompt", [], { - taskId: "test-task-id", - tools: testTools, + const deltaChunks = chunks.filter((chunk) => chunk.type === "tool_call_delta") + expect(deltaChunks).toHaveLength(1) + expect(deltaChunks[0]).toEqual({ + type: "tool_call_delta", + id: "call_lmstudio_123", + delta: '{"arg1":"value"}', }) - const chunks = [] - for await (const chunk of stream) { - // Simulate what Task.ts does: when we receive tool_call_partial, - // process it through NativeToolCallParser to populate rawChunkTracker - if (chunk.type === "tool_call_partial") { - NativeToolCallParser.processRawChunk({ - index: chunk.index, - id: chunk.id, - name: chunk.name, - arguments: chunk.arguments, - }) - } - chunks.push(chunk) - } - - // Should have tool_call_partial and tool_call_end - const partialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial") const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end") - - expect(partialChunks).toHaveLength(1) expect(endChunks).toHaveLength(1) - expect(endChunks[0].id).toBe("call_lmstudio_test") - }) - - it("should work with parallel tool calls disabled (sends false)", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Response" } }], - } - }, - })) - - const stream = handler.createMessage("test prompt", [], { - taskId: "test-task-id", - tools: testTools, - parallelToolCalls: false, + expect(endChunks[0]).toEqual({ + type: "tool_call_end", + id: "call_lmstudio_123", }) - await stream.next() - - // When parallelToolCalls is false, the parameter should be sent as false - const callArgs = mockCreate.mock.calls[0][0] - expect(callArgs).toHaveProperty("parallel_tool_calls", false) }) it("should handle reasoning content alongside tool calls", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [ - { - delta: { - content: "Thinking about this...", - }, - }, - ], - } - yield { - choices: [ - { - delta: { - tool_calls: [ - { - index: 0, - id: "call_after_think", - function: { - name: "test_tool", - arguments: '{"arg1":"result"}', - }, - }, - ], - }, - }, - ], - } - yield { - choices: [ - { - delta: {}, - finish_reason: "tool_calls", - }, - ], - } - }, - })) + async function* mockFullStream() { + yield { + type: "reasoning", + text: "Thinking about this...", + } + yield { + type: "tool-input-start", + id: "call_after_think", + toolName: "test_tool", + } + yield { + type: "tool-input-delta", + id: "call_after_think", + delta: '{"arg1":"result"}', + } + yield { + type: "tool-input-end", + id: "call_after_think", + } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }), + }) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -352,25 +194,60 @@ describe("LmStudioHandler Native Tools", () => { const chunks = [] for await (const chunk of stream) { - if (chunk.type === "tool_call_partial") { - NativeToolCallParser.processRawChunk({ - index: chunk.index, - id: chunk.id, - name: chunk.name, - arguments: chunk.arguments, - }) - } chunks.push(chunk) } - // Should have reasoning, tool_call_partial, and tool_call_end const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning") - const partialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial") + const startChunks = chunks.filter((chunk) => chunk.type === "tool_call_start") const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end") expect(reasoningChunks).toHaveLength(1) expect(reasoningChunks[0].text).toBe("Thinking about this...") - expect(partialChunks).toHaveLength(1) + expect(startChunks).toHaveLength(1) + expect(endChunks).toHaveLength(1) + }) + + it("should handle text and tool calls in the same response", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Here's the result: " } + yield { + type: "tool-input-start", + id: "call_mixed", + toolName: "test_tool", + } + yield { + type: "tool-input-delta", + id: "call_mixed", + delta: '{"arg1":"mixed"}', + } + yield { + type: "tool-input-end", + id: "call_mixed", + } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }), + }) + + const stream = handler.createMessage("test prompt", [], { + taskId: "test-task-id", + tools: testTools, + }) + + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const textChunks = chunks.filter((chunk) => chunk.type === "text") + const startChunks = chunks.filter((chunk) => chunk.type === "tool_call_start") + const endChunks = chunks.filter((chunk) => chunk.type === "tool_call_end") + + expect(textChunks).toHaveLength(1) + expect(textChunks[0].text).toBe("Here's the result: ") + expect(startChunks).toHaveLength(1) expect(endChunks).toHaveLength(1) }) }) diff --git a/src/api/providers/__tests__/lmstudio.spec.ts b/src/api/providers/__tests__/lmstudio.spec.ts index 0adebdeea7..0f6944e8da 100644 --- a/src/api/providers/__tests__/lmstudio.spec.ts +++ b/src/api/providers/__tests__/lmstudio.spec.ts @@ -1,63 +1,29 @@ -// Mock OpenAI client - must come before other imports -const mockCreate = vi.fn() -vi.mock("openai", () => { - return { - __esModule: true, - default: vi.fn().mockImplementation(() => ({ - chat: { - completions: { - create: mockCreate.mockImplementation(async (options) => { - if (!options.stream) { - return { - id: "test-completion", - choices: [ - { - message: { role: "assistant", content: "Test response" }, - finish_reason: "stop", - index: 0, - }, - ], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - total_tokens: 15, - }, - } - } +// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls +const { mockStreamText, mockGenerateText, mockWrapLanguageModel } = vi.hoisted(() => ({ + mockStreamText: vi.fn(), + mockGenerateText: vi.fn(), + mockWrapLanguageModel: vi.fn((opts: any) => opts.model), +})) - return { - [Symbol.asyncIterator]: async function* () { - yield { - choices: [ - { - delta: { content: "Test response" }, - index: 0, - }, - ], - usage: null, - } - yield { - choices: [ - { - delta: {}, - index: 0, - }, - ], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - total_tokens: 15, - }, - } - }, - } - }), - }, - }, - })), +vi.mock("ai", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + streamText: mockStreamText, + generateText: mockGenerateText, + wrapLanguageModel: mockWrapLanguageModel, } }) +vi.mock("@ai-sdk/openai-compatible", () => ({ + createOpenAICompatible: vi.fn(() => { + return vi.fn(() => ({ + modelId: "local-model", + provider: "lmstudio", + })) + }), +})) + import type { Anthropic } from "@anthropic-ai/sdk" import { LmStudioHandler } from "../lm-studio" @@ -74,7 +40,7 @@ describe("LmStudioHandler", () => { lmStudioBaseUrl: "http://localhost:1234", } handler = new LmStudioHandler(mockOptions) - mockCreate.mockClear() + vi.clearAllMocks() }) describe("constructor", () => { @@ -102,6 +68,20 @@ describe("LmStudioHandler", () => { ] it("should handle streaming responses", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test response" } + } + + const mockUsage = Promise.resolve({ + inputTokens: 10, + outputTokens: 5, + }) + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: mockUsage, + }) + const stream = handler.createMessage(systemPrompt, messages) const chunks: any[] = [] for await (const chunk of stream) { @@ -114,8 +94,43 @@ describe("LmStudioHandler", () => { expect(textChunks[0].text).toBe("Test response") }) + it("should include usage information", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test response" } + } + + const mockUsage = Promise.resolve({ + inputTokens: 10, + outputTokens: 5, + }) + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: mockUsage, + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunks = chunks.filter((chunk) => chunk.type === "usage") + expect(usageChunks.length).toBeGreaterThan(0) + expect(usageChunks[0].inputTokens).toBe(10) + expect(usageChunks[0].outputTokens).toBe(5) + }) + it("should handle API errors", async () => { - mockCreate.mockRejectedValueOnce(new Error("API Error")) + async function* mockFullStream(): AsyncGenerator<{ type: string; text: string }> { + yield { type: "text-delta", text: "" } + throw new Error("API Error") + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }), + }) const stream = handler.createMessage(systemPrompt, messages) @@ -123,36 +138,37 @@ describe("LmStudioHandler", () => { for await (const _chunk of stream) { // Should not reach here } - }).rejects.toThrow("Please check the LM Studio developer logs to debug what went wrong") + }).rejects.toThrow() }) }) describe("completePrompt", () => { it("should complete prompt successfully", async () => { + mockGenerateText.mockResolvedValue({ + text: "Test response", + }) + const result = await handler.completePrompt("Test prompt") expect(result).toBe("Test response") - expect(mockCreate).toHaveBeenCalledWith({ - model: mockOptions.lmStudioModelId, - messages: [{ role: "user", content: "Test prompt" }], - temperature: 0, - stream: false, - }) - }) - - it("should handle API errors", async () => { - mockCreate.mockRejectedValueOnce(new Error("API Error")) - await expect(handler.completePrompt("Test prompt")).rejects.toThrow( - "Please check the LM Studio developer logs to debug what went wrong", + expect(mockGenerateText).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: "Test prompt", + }), ) }) it("should handle empty response", async () => { - mockCreate.mockResolvedValueOnce({ - choices: [{ message: { content: "" } }], + mockGenerateText.mockResolvedValue({ + text: "", }) const result = await handler.completePrompt("Test prompt") expect(result).toBe("") }) + + it("should handle API errors with handleAiSdkError", async () => { + mockGenerateText.mockRejectedValueOnce(new Error("Connection refused")) + await expect(handler.completePrompt("Test prompt")).rejects.toThrow("LM Studio") + }) }) describe("getModel", () => { @@ -164,4 +180,131 @@ describe("LmStudioHandler", () => { expect(modelInfo.info.contextWindow).toBe(128_000) }) }) + + describe("speculative decoding", () => { + it("should include draft_model in providerOptions when speculative decoding is enabled", async () => { + const speculativeHandler = new LmStudioHandler({ + ...mockOptions, + lmStudioSpeculativeDecodingEnabled: true, + lmStudioDraftModelId: "draft-model-id", + }) + + async function* mockFullStream() { + yield { type: "text-delta", text: "Test" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 5, outputTokens: 3 }), + }) + + const stream = speculativeHandler.createMessage("test prompt", []) + for await (const _chunk of stream) { + // consume stream + } + + expect(mockStreamText).toHaveBeenCalledWith( + expect.objectContaining({ + providerOptions: { + lmstudio: { draft_model: "draft-model-id" }, + }, + }), + ) + }) + + it("should not include draft_model when speculative decoding is disabled", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 5, outputTokens: 3 }), + }) + + const stream = handler.createMessage("test prompt", []) + for await (const _chunk of stream) { + // consume stream + } + + const callArgs = mockStreamText.mock.calls[0][0] + expect(callArgs.providerOptions).toBeUndefined() + }) + + it("should include draft_model in completePrompt when speculative decoding is enabled", async () => { + const speculativeHandler = new LmStudioHandler({ + ...mockOptions, + lmStudioSpeculativeDecodingEnabled: true, + lmStudioDraftModelId: "draft-model-id", + }) + + mockGenerateText.mockResolvedValue({ text: "Test" }) + + await speculativeHandler.completePrompt("Test prompt") + + expect(mockGenerateText).toHaveBeenCalledWith( + expect.objectContaining({ + providerOptions: { + lmstudio: { draft_model: "draft-model-id" }, + }, + }), + ) + }) + }) + + describe("reasoning middleware", () => { + it("should wrap the language model with extractReasoningMiddleware for tags", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 5, outputTokens: 3 }), + }) + + const stream = handler.createMessage("test prompt", []) + for await (const _chunk of stream) { + // consume stream to trigger getLanguageModel() + } + + expect(mockWrapLanguageModel).toHaveBeenCalledWith( + expect.objectContaining({ + middleware: expect.any(Object), + }), + ) + }) + + it("should handle reasoning-delta chunks from middleware-processed stream", async () => { + async function* mockFullStream() { + yield { type: "reasoning-delta", text: "Let me think about this..." } + yield { type: "text-delta", text: "The answer is 42." } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ inputTokens: 10, outputTokens: 8 }), + }) + + const stream = handler.createMessage("test prompt", []) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const reasoningChunks = chunks.filter((c) => c.type === "reasoning") + const textChunks = chunks.filter((c) => c.type === "text") + + expect(reasoningChunks).toHaveLength(1) + expect(reasoningChunks[0].text).toBe("Let me think about this...") + expect(textChunks).toHaveLength(1) + expect(textChunks[0].text).toBe("The answer is 42.") + }) + }) + + describe("isAiSdkProvider", () => { + it("should return true", () => { + expect(handler.isAiSdkProvider()).toBe(true) + }) + }) }) diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index a771394c53..fdc95afb41 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -1,39 +1,49 @@ import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" -import axios from "axios" +import { streamText, generateText, ToolSet, wrapLanguageModel, extractReasoningMiddleware, LanguageModel } from "ai" import { type ModelInfo, openAiModelInfoSaneDefaults, LMSTUDIO_DEFAULT_TEMPERATURE } from "@roo-code/types" import type { ApiHandlerOptions } from "../../shared/api" -import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser" -import { TagMatcher } from "../../utils/tag-matcher" - -import { convertToOpenAiMessages } from "../transform/openai-format" +import { + convertToAiSdkMessages, + convertToolsForAiSdk, + processAiSdkStreamPart, + mapToolChoice, + handleAiSdkError, +} from "../transform/ai-sdk" import { ApiStream } from "../transform/stream" -import { BaseProvider } from "./base-provider" +import { OpenAICompatibleHandler, OpenAICompatibleConfig } from "./openai-compatible" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import { getModelsFromCache } from "./fetchers/modelCache" -import { getApiRequestTimeout } from "./utils/timeout-config" -import { handleOpenAIError } from "./utils/openai-error-handler" - -export class LmStudioHandler extends BaseProvider implements SingleCompletionHandler { - protected options: ApiHandlerOptions - private client: OpenAI - private readonly providerName = "LM Studio" +export class LmStudioHandler extends OpenAICompatibleHandler implements SingleCompletionHandler { constructor(options: ApiHandlerOptions) { - super() - this.options = options + const modelId = options.lmStudioModelId || "" + const baseURL = (options.lmStudioBaseUrl || "http://localhost:1234") + "/v1" - // LM Studio uses "noop" as a placeholder API key - const apiKey = "noop" + const models = getModelsFromCache("lmstudio") + const modelInfo = (models && modelId && models[modelId]) || openAiModelInfoSaneDefaults - this.client = new OpenAI({ - baseURL: (this.options.lmStudioBaseUrl || "http://localhost:1234") + "/v1", - apiKey: apiKey, - timeout: getApiRequestTimeout(), + const config: OpenAICompatibleConfig = { + providerName: "lmstudio", + baseURL, + apiKey: "noop", + modelId, + modelInfo, + temperature: options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, + modelMaxTokens: options.modelMaxTokens ?? undefined, + } + + super(options, config) + } + + protected override getLanguageModel(): LanguageModel { + const baseModel = this.provider(this.config.modelId) + return wrapLanguageModel({ + model: baseModel, + middleware: extractReasoningMiddleware({ tagName: "think" }), }) } @@ -42,189 +52,83 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ - { role: "system", content: systemPrompt }, - ...convertToOpenAiMessages(messages), - ] + const model = this.getModel() + const languageModel = this.getLanguageModel() - // ------------------------- - // Track token usage - // ------------------------- - const toContentBlocks = ( - blocks: Anthropic.Messages.MessageParam[] | string, - ): Anthropic.Messages.ContentBlockParam[] => { - if (typeof blocks === "string") { - return [{ type: "text", text: blocks }] - } + const aiSdkMessages = convertToAiSdkMessages(messages) - const result: Anthropic.Messages.ContentBlockParam[] = [] - for (const msg of blocks) { - if (typeof msg.content === "string") { - result.push({ type: "text", text: msg.content }) - } else if (Array.isArray(msg.content)) { - for (const part of msg.content) { - if (part.type === "text") { - result.push({ type: "text", text: part.text }) - } - } - } - } - return result + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) + const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + + const requestOptions: Parameters[0] = { + model: languageModel, + system: systemPrompt, + messages: aiSdkMessages, + temperature: model.temperature ?? this.config.temperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, + maxOutputTokens: this.getMaxOutputTokens(), + tools: aiSdkTools, + toolChoice: mapToolChoice(metadata?.tool_choice), } - let inputTokens = 0 - try { - inputTokens = await this.countTokens([{ type: "text", text: systemPrompt }, ...toContentBlocks(messages)]) - } catch (err) { - console.error("[LmStudio] Failed to count input tokens:", err) - inputTokens = 0 + if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) { + requestOptions.providerOptions = { + lmstudio: { draft_model: this.options.lmStudioDraftModelId }, + } } - let assistantText = "" + const result = streamText(requestOptions) try { - const params: OpenAI.Chat.ChatCompletionCreateParamsStreaming & { draft_model?: string } = { - model: this.getModel().id, - messages: openAiMessages, - temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, - stream: true, - tools: this.convertToolsForOpenAI(metadata?.tools), - tool_choice: metadata?.tool_choice, - parallel_tool_calls: metadata?.parallelToolCalls ?? true, - } - - if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) { - params.draft_model = this.options.lmStudioDraftModelId - } - - let results - try { - results = await this.client.chat.completions.create(params) - } catch (error) { - throw handleOpenAIError(error, this.providerName) - } - - const matcher = new TagMatcher( - "think", - (chunk) => - ({ - type: chunk.matched ? "reasoning" : "text", - text: chunk.data, - }) as const, - ) - - for await (const chunk of results) { - const delta = chunk.choices[0]?.delta - const finishReason = chunk.choices[0]?.finish_reason - - if (delta?.content) { - assistantText += delta.content - for (const processedChunk of matcher.update(delta.content)) { - yield processedChunk - } - } - - // Handle tool calls in stream - emit partial chunks for NativeToolCallParser - if (delta?.tool_calls) { - for (const toolCall of delta.tool_calls) { - yield { - type: "tool_call_partial", - index: toolCall.index, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, - } - } - } - - // Process finish_reason to emit tool_call_end events - if (finishReason) { - const endEvents = NativeToolCallParser.processFinishReason(finishReason) - for (const event of endEvents) { - yield event - } + for await (const part of result.fullStream) { + for (const chunk of processAiSdkStreamPart(part)) { + yield chunk } } - for (const processedChunk of matcher.final()) { - yield processedChunk + const usage = await result.usage + if (usage) { + yield this.processUsageMetrics(usage) } - - let outputTokens = 0 - try { - outputTokens = await this.countTokens([{ type: "text", text: assistantText }]) - } catch (err) { - console.error("[LmStudio] Failed to count output tokens:", err) - outputTokens = 0 - } - - yield { - type: "usage", - inputTokens, - outputTokens, - } as const } catch (error) { - throw new Error( - "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Roo Code's prompts.", - ) + throw handleAiSdkError(error, "LM Studio") } } - override getModel(): { id: string; info: ModelInfo } { + override getModel(): { id: string; info: ModelInfo; maxTokens?: number; temperature?: number } { const models = getModelsFromCache("lmstudio") - if (models && this.options.lmStudioModelId && models[this.options.lmStudioModelId]) { - return { - id: this.options.lmStudioModelId, - info: models[this.options.lmStudioModelId], - } - } else { - return { - id: this.options.lmStudioModelId || "", - info: openAiModelInfoSaneDefaults, - } + const modelId = this.options.lmStudioModelId || "" + + const info = (models && modelId && models[modelId]) || openAiModelInfoSaneDefaults + + return { + id: modelId, + info, + temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, + maxTokens: this.options.modelMaxTokens ?? undefined, } } - async completePrompt(prompt: string): Promise { + override async completePrompt(prompt: string): Promise { + const languageModel = this.getLanguageModel() + + const options: Parameters[0] = { + model: languageModel, + prompt, + maxOutputTokens: this.getMaxOutputTokens(), + temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, + } + + if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) { + options.providerOptions = { + lmstudio: { draft_model: this.options.lmStudioDraftModelId }, + } + } + try { - // Create params object with optional draft model - const params: any = { - model: this.getModel().id, - messages: [{ role: "user", content: prompt }], - temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, - stream: false, - } - - // Add draft model if speculative decoding is enabled and a draft model is specified - if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) { - params.draft_model = this.options.lmStudioDraftModelId - } - - let response - try { - response = await this.client.chat.completions.create(params) - } catch (error) { - throw handleOpenAIError(error, this.providerName) - } - return response.choices[0]?.message.content || "" + const { text } = await generateText(options) + return text } catch (error) { - throw new Error( - "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Roo Code's prompts.", - ) + throw handleAiSdkError(error, "LM Studio") } } } - -export async function getLmStudioModels(baseUrl = "http://localhost:1234") { - try { - if (!URL.canParse(baseUrl)) { - return [] - } - - const response = await axios.get(`${baseUrl}/v1/models`) - const modelsArray = response.data?.data?.map((model: any) => model.id) || [] - return [...new Set(modelsArray)] - } catch (error) { - return [] - } -}