mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
feat: Add OpenAI-compatible base provider and migrate Moonshot to AI SDK (#11063)
This commit is contained in:
parent
ed35b09aad
commit
8d38e60187
7 changed files with 552 additions and 154 deletions
|
|
@ -53,6 +53,19 @@ export const moonshotModels = {
|
|||
defaultTemperature: 1.0,
|
||||
description: `The kimi-k2-thinking model is a general-purpose agentic reasoning model developed by Moonshot AI. Thanks to its strength in deep reasoning and multi-turn tool use, it can solve even the hardest problems.`,
|
||||
},
|
||||
"kimi-k2.5": {
|
||||
maxTokens: 16_384,
|
||||
contextWindow: 262_144,
|
||||
supportsImages: false,
|
||||
supportsPromptCache: true,
|
||||
inputPrice: 0.6, // $0.60 per million tokens (cache miss)
|
||||
outputPrice: 3.0, // $3.00 per million tokens
|
||||
cacheReadsPrice: 0.1, // $0.10 per million tokens (cache hit)
|
||||
supportsTemperature: true,
|
||||
defaultTemperature: 1.0,
|
||||
description:
|
||||
"Kimi K2.5 is the latest generation of Moonshot AI's Kimi series, featuring improved reasoning capabilities and enhanced performance across diverse tasks.",
|
||||
},
|
||||
} as const satisfies Record<string, ModelInfo>
|
||||
|
||||
export const MOONSHOT_DEFAULT_TEMPERATURE = 0.6
|
||||
|
|
|
|||
36
pnpm-lock.yaml
generated
36
pnpm-lock.yaml
generated
|
|
@ -990,6 +990,9 @@ importers:
|
|||
specifier: 3.25.76
|
||||
version: 3.25.76
|
||||
devDependencies:
|
||||
'@ai-sdk/openai-compatible':
|
||||
specifier: ^1.0.0
|
||||
version: 1.0.31(zod@3.25.76)
|
||||
'@openrouter/ai-sdk-provider':
|
||||
specifier: ^2.0.4
|
||||
version: 2.1.1(ai@6.0.57(zod@3.25.76))(zod@3.25.76)
|
||||
|
|
@ -1384,12 +1387,28 @@ packages:
|
|||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/openai-compatible@1.0.31':
|
||||
resolution: {integrity: sha512-znBvaVHM0M6yWNerIEy3hR+O8ZK2sPcE7e2cxfb6kYLEX3k//JH5VDnRnajseVofg7LXtTCFFdjsB7WLf1BdeQ==}
|
||||
engines: {node: '>=18'}
|
||||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider-utils@3.0.20':
|
||||
resolution: {integrity: sha512-iXHVe0apM2zUEzauqJwqmpC37A5rihrStAih5Ks+JE32iTe4LZ58y17UGBjpQQTCRw9YxMeo2UFLxLpBluyvLQ==}
|
||||
engines: {node: '>=18'}
|
||||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider-utils@4.0.10':
|
||||
resolution: {integrity: sha512-VeDAiCH+ZK8Xs4hb9Cw7pHlujWNL52RKe8TExOkrw6Ir1AmfajBZTb9XUdKOZO08RwQElIKA8+Ltm+Gqfo8djQ==}
|
||||
engines: {node: '>=18'}
|
||||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider@2.0.1':
|
||||
resolution: {integrity: sha512-KCUwswvsC5VsW2PWFqF8eJgSCu5Ysj7m1TxiHTVA6g7k360bk0RNQENT8KTMAYEs+8fWPD3Uu4dEmzGHc+jGng==}
|
||||
engines: {node: '>=18'}
|
||||
|
||||
'@ai-sdk/provider@3.0.5':
|
||||
resolution: {integrity: sha512-2Xmoq6DBJqmSl80U6V9z5jJSJP7ehaJJQMy2iFUqTay06wdCqTnPVBBQbtEL8RCChenL+q5DC5H5WzU3vV3v8w==}
|
||||
engines: {node: '>=18'}
|
||||
|
|
@ -10784,6 +10803,19 @@ snapshots:
|
|||
'@vercel/oidc': 3.1.0
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/openai-compatible@1.0.31(zod@3.25.76)':
|
||||
dependencies:
|
||||
'@ai-sdk/provider': 2.0.1
|
||||
'@ai-sdk/provider-utils': 3.0.20(zod@3.25.76)
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider-utils@3.0.20(zod@3.25.76)':
|
||||
dependencies:
|
||||
'@ai-sdk/provider': 2.0.1
|
||||
'@standard-schema/spec': 1.1.0
|
||||
eventsource-parser: 3.0.6
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider-utils@4.0.10(zod@3.25.76)':
|
||||
dependencies:
|
||||
'@ai-sdk/provider': 3.0.5
|
||||
|
|
@ -10791,6 +10823,10 @@ snapshots:
|
|||
eventsource-parser: 3.0.6
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider@2.0.1':
|
||||
dependencies:
|
||||
json-schema: 0.4.0
|
||||
|
||||
'@ai-sdk/provider@3.0.5':
|
||||
dependencies:
|
||||
json-schema: 0.4.0
|
||||
|
|
|
|||
|
|
@ -1,67 +1,28 @@
|
|||
// Mocks must come first, before 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", refusal: null },
|
||||
finish_reason: "stop",
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 5,
|
||||
total_tokens: 15,
|
||||
cached_tokens: 2,
|
||||
},
|
||||
}
|
||||
}
|
||||
// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
|
||||
const { mockStreamText, mockGenerateText } = vi.hoisted(() => ({
|
||||
mockStreamText: vi.fn(),
|
||||
mockGenerateText: vi.fn(),
|
||||
}))
|
||||
|
||||
// Return async iterator for streaming
|
||||
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,
|
||||
cached_tokens: 2,
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
}),
|
||||
},
|
||||
},
|
||||
})),
|
||||
vi.mock("ai", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("ai")>()
|
||||
return {
|
||||
...actual,
|
||||
streamText: mockStreamText,
|
||||
generateText: mockGenerateText,
|
||||
}
|
||||
})
|
||||
|
||||
import OpenAI from "openai"
|
||||
vi.mock("@ai-sdk/openai-compatible", () => ({
|
||||
createOpenAICompatible: vi.fn(() => {
|
||||
// Return a function that returns a mock language model
|
||||
return vi.fn(() => ({
|
||||
modelId: "moonshot-chat",
|
||||
provider: "moonshot",
|
||||
}))
|
||||
}),
|
||||
}))
|
||||
|
||||
import type { Anthropic } from "@anthropic-ai/sdk"
|
||||
|
||||
import { moonshotDefaultModelId } from "@roo-code/types"
|
||||
|
|
@ -90,15 +51,6 @@ describe("MoonshotHandler", () => {
|
|||
expect(handler.getModel().id).toBe(mockOptions.apiModelId)
|
||||
})
|
||||
|
||||
it.skip("should throw error if API key is missing", () => {
|
||||
expect(() => {
|
||||
new MoonshotHandler({
|
||||
...mockOptions,
|
||||
moonshotApiKey: undefined,
|
||||
})
|
||||
}).toThrow("Moonshot API key is required")
|
||||
})
|
||||
|
||||
it("should use default model ID if not provided", () => {
|
||||
const handlerWithoutModel = new MoonshotHandler({
|
||||
...mockOptions,
|
||||
|
|
@ -113,12 +65,6 @@ describe("MoonshotHandler", () => {
|
|||
moonshotBaseUrl: undefined,
|
||||
})
|
||||
expect(handlerWithoutBaseUrl).toBeInstanceOf(MoonshotHandler)
|
||||
// The base URL is passed to OpenAI client internally
|
||||
expect(OpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: "https://api.moonshot.ai/v1",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should use chinese base URL if provided", () => {
|
||||
|
|
@ -128,18 +74,6 @@ describe("MoonshotHandler", () => {
|
|||
moonshotBaseUrl: customBaseUrl,
|
||||
})
|
||||
expect(handlerWithCustomUrl).toBeInstanceOf(MoonshotHandler)
|
||||
// The custom base URL is passed to OpenAI client
|
||||
expect(OpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: customBaseUrl,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should set includeMaxTokens to true", () => {
|
||||
// Create a new handler and verify OpenAI client was called with includeMaxTokens
|
||||
const _handler = new MoonshotHandler(mockOptions)
|
||||
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: mockOptions.moonshotApiKey }))
|
||||
})
|
||||
})
|
||||
|
||||
|
|
@ -151,7 +85,7 @@ describe("MoonshotHandler", () => {
|
|||
expect(model.info.maxTokens).toBe(16384)
|
||||
expect(model.info.contextWindow).toBe(262144)
|
||||
expect(model.info.supportsImages).toBe(false)
|
||||
expect(model.info.supportsPromptCache).toBe(true) // Should be true now
|
||||
expect(model.info.supportsPromptCache).toBe(true)
|
||||
})
|
||||
|
||||
it("should return provided model ID with default model info if model does not exist", () => {
|
||||
|
|
@ -162,11 +96,8 @@ describe("MoonshotHandler", () => {
|
|||
const model = handlerWithInvalidModel.getModel()
|
||||
expect(model.id).toBe("invalid-model") // Returns provided ID
|
||||
expect(model.info).toBeDefined()
|
||||
// With the current implementation, it's the same object reference when using default model info
|
||||
expect(model.info).toBe(handler.getModel().info)
|
||||
// Should have the same base properties
|
||||
// Should have the same base properties as default model
|
||||
expect(model.info.contextWindow).toBe(handler.getModel().info.contextWindow)
|
||||
// And should have supportsPromptCache set to true
|
||||
expect(model.info.supportsPromptCache).toBe(true)
|
||||
})
|
||||
|
||||
|
|
@ -203,6 +134,24 @@ describe("MoonshotHandler", () => {
|
|||
]
|
||||
|
||||
it("should handle streaming responses", async () => {
|
||||
// Mock the fullStream async generator
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
// Mock usage promise
|
||||
const mockUsage = Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
details: { cachedInputTokens: undefined },
|
||||
raw: { cached_tokens: 2 },
|
||||
})
|
||||
|
||||
mockStreamText.mockReturnValue({
|
||||
fullStream: mockFullStream(),
|
||||
usage: mockUsage,
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
|
|
@ -216,6 +165,22 @@ describe("MoonshotHandler", () => {
|
|||
})
|
||||
|
||||
it("should include usage information", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
const mockUsage = Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
details: {},
|
||||
raw: { cached_tokens: 2 },
|
||||
})
|
||||
|
||||
mockStreamText.mockReturnValue({
|
||||
fullStream: mockFullStream(),
|
||||
usage: mockUsage,
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
|
|
@ -229,6 +194,22 @@ describe("MoonshotHandler", () => {
|
|||
})
|
||||
|
||||
it("should include cache metrics in usage information", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
const mockUsage = Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
details: {},
|
||||
raw: { cached_tokens: 2 },
|
||||
})
|
||||
|
||||
mockStreamText.mockReturnValue({
|
||||
fullStream: mockFullStream(),
|
||||
usage: mockUsage,
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
|
|
@ -242,6 +223,23 @@ describe("MoonshotHandler", () => {
|
|||
})
|
||||
})
|
||||
|
||||
describe("completePrompt", () => {
|
||||
it("should complete a prompt using generateText", async () => {
|
||||
mockGenerateText.mockResolvedValue({
|
||||
text: "Test completion",
|
||||
})
|
||||
|
||||
const result = await handler.completePrompt("Test prompt")
|
||||
|
||||
expect(result).toBe("Test completion")
|
||||
expect(mockGenerateText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
prompt: "Test prompt",
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("processUsageMetrics", () => {
|
||||
it("should correctly process usage metrics including cache information", () => {
|
||||
// We need to access the protected method, so we'll create a test subclass
|
||||
|
|
@ -254,10 +252,12 @@ describe("MoonshotHandler", () => {
|
|||
const testHandler = new TestMoonshotHandler(mockOptions)
|
||||
|
||||
const usage = {
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 50,
|
||||
total_tokens: 150,
|
||||
cached_tokens: 20,
|
||||
inputTokens: 100,
|
||||
outputTokens: 50,
|
||||
details: {},
|
||||
raw: {
|
||||
cached_tokens: 20,
|
||||
},
|
||||
}
|
||||
|
||||
const result = testHandler.testProcessUsageMetrics(usage)
|
||||
|
|
@ -279,10 +279,10 @@ describe("MoonshotHandler", () => {
|
|||
const testHandler = new TestMoonshotHandler(mockOptions)
|
||||
|
||||
const usage = {
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 50,
|
||||
total_tokens: 150,
|
||||
// No cached_tokens
|
||||
inputTokens: 100,
|
||||
outputTokens: 50,
|
||||
details: {},
|
||||
raw: {},
|
||||
}
|
||||
|
||||
const result = testHandler.testProcessUsageMetrics(usage)
|
||||
|
|
@ -295,31 +295,25 @@ describe("MoonshotHandler", () => {
|
|||
})
|
||||
})
|
||||
|
||||
describe("addMaxTokensIfNeeded", () => {
|
||||
it("should always add max_tokens regardless of includeMaxTokens option", () => {
|
||||
// Create a test subclass to access the protected method
|
||||
describe("getMaxOutputTokens", () => {
|
||||
it("should return maxTokens from model info", () => {
|
||||
class TestMoonshotHandler extends MoonshotHandler {
|
||||
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
public testGetMaxOutputTokens() {
|
||||
return this.getMaxOutputTokens()
|
||||
}
|
||||
}
|
||||
|
||||
const testHandler = new TestMoonshotHandler(mockOptions)
|
||||
const requestOptions: any = {}
|
||||
const modelInfo = {
|
||||
maxTokens: 32_000,
|
||||
}
|
||||
const result = testHandler.testGetMaxOutputTokens()
|
||||
|
||||
// Test with includeMaxTokens set to false - should still add max tokens
|
||||
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
|
||||
expect(requestOptions.max_tokens).toBe(32_000)
|
||||
// Default model maxTokens is 16384
|
||||
expect(result).toBe(16384)
|
||||
})
|
||||
|
||||
it("should use modelMaxTokens when provided", () => {
|
||||
class TestMoonshotHandler extends MoonshotHandler {
|
||||
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
public testGetMaxOutputTokens() {
|
||||
return this.getMaxOutputTokens()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -328,32 +322,153 @@ describe("MoonshotHandler", () => {
|
|||
...mockOptions,
|
||||
modelMaxTokens: customMaxTokens,
|
||||
})
|
||||
const requestOptions: any = {}
|
||||
const modelInfo = {
|
||||
maxTokens: 32_000,
|
||||
}
|
||||
|
||||
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
|
||||
expect(requestOptions.max_tokens).toBe(customMaxTokens)
|
||||
const result = testHandler.testGetMaxOutputTokens()
|
||||
expect(result).toBe(customMaxTokens)
|
||||
})
|
||||
|
||||
it("should fall back to modelInfo.maxTokens when modelMaxTokens is not provided", () => {
|
||||
class TestMoonshotHandler extends MoonshotHandler {
|
||||
public testAddMaxTokensIfNeeded(requestOptions: any, modelInfo: any) {
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
public testGetMaxOutputTokens() {
|
||||
return this.getMaxOutputTokens()
|
||||
}
|
||||
}
|
||||
|
||||
const testHandler = new TestMoonshotHandler(mockOptions)
|
||||
const requestOptions: any = {}
|
||||
const modelInfo = {
|
||||
maxTokens: 16_000,
|
||||
const result = testHandler.testGetMaxOutputTokens()
|
||||
|
||||
// moonshot-chat has maxTokens of 16384
|
||||
expect(result).toBe(16384)
|
||||
})
|
||||
})
|
||||
|
||||
describe("tool handling", () => {
|
||||
const systemPrompt = "You are a helpful assistant."
|
||||
const messages: Anthropic.Messages.MessageParam[] = [
|
||||
{
|
||||
role: "user",
|
||||
content: [{ type: "text" as const, text: "Hello!" }],
|
||||
},
|
||||
]
|
||||
|
||||
it("should handle tool calls in streaming", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield {
|
||||
type: "tool-input-start",
|
||||
id: "tool-call-1",
|
||||
toolName: "read_file",
|
||||
}
|
||||
yield {
|
||||
type: "tool-input-delta",
|
||||
id: "tool-call-1",
|
||||
delta: '{"path":"test.ts"}',
|
||||
}
|
||||
yield {
|
||||
type: "tool-input-end",
|
||||
id: "tool-call-1",
|
||||
}
|
||||
}
|
||||
|
||||
testHandler.testAddMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
const mockUsage = Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
details: {},
|
||||
raw: {},
|
||||
})
|
||||
|
||||
expect(requestOptions.max_tokens).toBe(16_000)
|
||||
mockStreamText.mockReturnValue({
|
||||
fullStream: mockFullStream(),
|
||||
usage: mockUsage,
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages, {
|
||||
taskId: "test-task",
|
||||
tools: [
|
||||
{
|
||||
type: "function",
|
||||
function: {
|
||||
name: "read_file",
|
||||
description: "Read a file",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: { path: { type: "string" } },
|
||||
required: ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
const toolCallStartChunks = chunks.filter((c) => c.type === "tool_call_start")
|
||||
const toolCallDeltaChunks = chunks.filter((c) => c.type === "tool_call_delta")
|
||||
const toolCallEndChunks = chunks.filter((c) => c.type === "tool_call_end")
|
||||
|
||||
expect(toolCallStartChunks.length).toBe(1)
|
||||
expect(toolCallStartChunks[0].id).toBe("tool-call-1")
|
||||
expect(toolCallStartChunks[0].name).toBe("read_file")
|
||||
|
||||
expect(toolCallDeltaChunks.length).toBe(1)
|
||||
expect(toolCallDeltaChunks[0].delta).toBe('{"path":"test.ts"}')
|
||||
|
||||
expect(toolCallEndChunks.length).toBe(1)
|
||||
expect(toolCallEndChunks[0].id).toBe("tool-call-1")
|
||||
})
|
||||
|
||||
it("should handle complete tool calls", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield {
|
||||
type: "tool-call",
|
||||
toolCallId: "tool-call-1",
|
||||
toolName: "read_file",
|
||||
input: { path: "test.ts" },
|
||||
}
|
||||
}
|
||||
|
||||
const mockUsage = Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
details: {},
|
||||
raw: {},
|
||||
})
|
||||
|
||||
mockStreamText.mockReturnValue({
|
||||
fullStream: mockFullStream(),
|
||||
usage: mockUsage,
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages, {
|
||||
taskId: "test-task",
|
||||
tools: [
|
||||
{
|
||||
type: "function",
|
||||
function: {
|
||||
name: "read_file",
|
||||
description: "Read a file",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: { path: { type: "string" } },
|
||||
required: ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
const toolCallChunks = chunks.filter((c) => c.type === "tool_call")
|
||||
expect(toolCallChunks.length).toBe(1)
|
||||
expect(toolCallChunks[0].id).toBe("tool-call-1")
|
||||
expect(toolCallChunks[0].name).toBe("read_file")
|
||||
expect(toolCallChunks[0].arguments).toBe('{"path":"test.ts"}')
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ export { MistralHandler } from "./mistral"
|
|||
export { OpenAiCodexHandler } from "./openai-codex"
|
||||
export { OpenAiNativeHandler } from "./openai-native"
|
||||
export { OpenAiHandler } from "./openai"
|
||||
export { OpenAICompatibleHandler } from "./openai-compatible"
|
||||
export type { OpenAICompatibleConfig } from "./openai-compatible"
|
||||
export { OpenRouterHandler } from "./openrouter"
|
||||
export { QwenCodeHandler } from "./qwen-code"
|
||||
export { RequestyHandler } from "./requesty"
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import OpenAI from "openai"
|
||||
import { moonshotModels, moonshotDefaultModelId, type ModelInfo } from "@roo-code/types"
|
||||
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
|
@ -6,18 +5,25 @@ import type { ApiHandlerOptions } from "../../shared/api"
|
|||
import type { ApiStreamUsageChunk } from "../transform/stream"
|
||||
import { getModelParams } from "../transform/model-params"
|
||||
|
||||
import { OpenAiHandler } from "./openai"
|
||||
import { OpenAICompatibleHandler, OpenAICompatibleConfig } from "./openai-compatible"
|
||||
|
||||
export class MoonshotHandler extends OpenAiHandler {
|
||||
export class MoonshotHandler extends OpenAICompatibleHandler {
|
||||
constructor(options: ApiHandlerOptions) {
|
||||
super({
|
||||
...options,
|
||||
openAiApiKey: options.moonshotApiKey ?? "not-provided",
|
||||
openAiModelId: options.apiModelId ?? moonshotDefaultModelId,
|
||||
openAiBaseUrl: options.moonshotBaseUrl ?? "https://api.moonshot.ai/v1",
|
||||
openAiStreamingEnabled: true,
|
||||
includeMaxTokens: true,
|
||||
})
|
||||
const modelId = options.apiModelId ?? moonshotDefaultModelId
|
||||
const modelInfo =
|
||||
moonshotModels[modelId as keyof typeof moonshotModels] || moonshotModels[moonshotDefaultModelId]
|
||||
|
||||
const config: OpenAICompatibleConfig = {
|
||||
providerName: "moonshot",
|
||||
baseURL: options.moonshotBaseUrl ?? "https://api.moonshot.ai/v1",
|
||||
apiKey: options.moonshotApiKey ?? "not-provided",
|
||||
modelId,
|
||||
modelInfo,
|
||||
modelMaxTokens: options.modelMaxTokens ?? undefined,
|
||||
temperature: options.modelTemperature ?? undefined,
|
||||
}
|
||||
|
||||
super(options, config)
|
||||
}
|
||||
|
||||
override getModel() {
|
||||
|
|
@ -27,25 +33,38 @@ export class MoonshotHandler extends OpenAiHandler {
|
|||
return { id, info, ...params }
|
||||
}
|
||||
|
||||
// Override to handle Moonshot's usage metrics, including caching.
|
||||
protected override processUsageMetrics(usage: any): ApiStreamUsageChunk {
|
||||
/**
|
||||
* Override to handle Moonshot's usage metrics, including caching.
|
||||
* Moonshot returns cached_tokens in a different location than standard OpenAI.
|
||||
*/
|
||||
protected override processUsageMetrics(usage: {
|
||||
inputTokens?: number
|
||||
outputTokens?: number
|
||||
details?: {
|
||||
cachedInputTokens?: number
|
||||
reasoningTokens?: number
|
||||
}
|
||||
raw?: Record<string, unknown>
|
||||
}): ApiStreamUsageChunk {
|
||||
// Moonshot uses cached_tokens at the top level of raw usage data
|
||||
const rawUsage = usage.raw as { cached_tokens?: number } | undefined
|
||||
|
||||
return {
|
||||
type: "usage",
|
||||
inputTokens: usage?.prompt_tokens || 0,
|
||||
outputTokens: usage?.completion_tokens || 0,
|
||||
inputTokens: usage.inputTokens || 0,
|
||||
outputTokens: usage.outputTokens || 0,
|
||||
cacheWriteTokens: 0,
|
||||
cacheReadTokens: usage?.cached_tokens,
|
||||
cacheReadTokens: rawUsage?.cached_tokens ?? usage.details?.cachedInputTokens,
|
||||
}
|
||||
}
|
||||
|
||||
// Override to always include max_tokens for Moonshot (not max_completion_tokens)
|
||||
protected override addMaxTokensIfNeeded(
|
||||
requestOptions:
|
||||
| OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
|
||||
| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
|
||||
modelInfo: ModelInfo,
|
||||
): void {
|
||||
// Moonshot uses max_tokens instead of max_completion_tokens
|
||||
requestOptions.max_tokens = this.options.modelMaxTokens || modelInfo.maxTokens
|
||||
/**
|
||||
* Override to always include max_tokens for Moonshot (not max_completion_tokens).
|
||||
* Moonshot requires max_tokens parameter to be sent.
|
||||
*/
|
||||
protected override getMaxOutputTokens(): number | undefined {
|
||||
const modelInfo = this.config.modelInfo
|
||||
// Moonshot always requires max_tokens
|
||||
return this.options.modelMaxTokens || modelInfo.maxTokens || undefined
|
||||
}
|
||||
}
|
||||
|
|
|
|||
212
src/api/providers/openai-compatible.ts
Normal file
212
src/api/providers/openai-compatible.ts
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
/**
|
||||
* OpenAI-compatible provider base class using Vercel AI SDK.
|
||||
* This provides a parallel implementation to OpenAiHandler using @ai-sdk/openai-compatible.
|
||||
*/
|
||||
|
||||
import { Anthropic } from "@anthropic-ai/sdk"
|
||||
import OpenAI from "openai"
|
||||
import { createOpenAICompatible } from "@ai-sdk/openai-compatible"
|
||||
import { streamText, generateText, LanguageModel, ToolSet } from "ai"
|
||||
|
||||
import type { ModelInfo } from "@roo-code/types"
|
||||
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
||||
import { convertToAiSdkMessages, convertToolsForAiSdk, processAiSdkStreamPart } from "../transform/ai-sdk"
|
||||
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
|
||||
|
||||
import { DEFAULT_HEADERS } from "./constants"
|
||||
import { BaseProvider } from "./base-provider"
|
||||
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
|
||||
|
||||
/**
|
||||
* Configuration options for creating an OpenAI-compatible provider.
|
||||
*/
|
||||
export interface OpenAICompatibleConfig {
|
||||
/** Provider name for identification */
|
||||
providerName: string
|
||||
/** Base URL for the API endpoint */
|
||||
baseURL: string
|
||||
/** API key for authentication */
|
||||
apiKey: string
|
||||
/** Model ID to use */
|
||||
modelId: string
|
||||
/** Model information */
|
||||
modelInfo: ModelInfo
|
||||
/** Optional custom headers */
|
||||
headers?: Record<string, string>
|
||||
/** Whether to include max_tokens in requests (default: false uses max_completion_tokens) */
|
||||
useMaxTokens?: boolean
|
||||
/** User-configured max tokens override */
|
||||
modelMaxTokens?: number
|
||||
/** Temperature setting */
|
||||
temperature?: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Base class for OpenAI-compatible API providers using Vercel AI SDK.
|
||||
* Extends BaseProvider and implements SingleCompletionHandler.
|
||||
*/
|
||||
export abstract class OpenAICompatibleHandler extends BaseProvider implements SingleCompletionHandler {
|
||||
protected options: ApiHandlerOptions
|
||||
protected config: OpenAICompatibleConfig
|
||||
protected provider: ReturnType<typeof createOpenAICompatible>
|
||||
|
||||
constructor(options: ApiHandlerOptions, config: OpenAICompatibleConfig) {
|
||||
super()
|
||||
this.options = options
|
||||
this.config = config
|
||||
|
||||
// Create the OpenAI-compatible provider using AI SDK
|
||||
this.provider = createOpenAICompatible({
|
||||
name: config.providerName,
|
||||
baseURL: config.baseURL,
|
||||
apiKey: config.apiKey,
|
||||
headers: {
|
||||
...DEFAULT_HEADERS,
|
||||
...(config.headers || {}),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the language model for the configured model ID.
|
||||
*/
|
||||
protected getLanguageModel(): LanguageModel {
|
||||
return this.provider(this.config.modelId)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the model information. Must be implemented by subclasses.
|
||||
*/
|
||||
abstract override getModel(): { id: string; info: ModelInfo; maxTokens?: number; temperature?: number }
|
||||
|
||||
/**
|
||||
* Process usage metrics from the AI SDK response.
|
||||
* Can be overridden by subclasses to handle provider-specific usage formats.
|
||||
*/
|
||||
protected processUsageMetrics(usage: {
|
||||
inputTokens?: number
|
||||
outputTokens?: number
|
||||
details?: {
|
||||
cachedInputTokens?: number
|
||||
reasoningTokens?: number
|
||||
}
|
||||
raw?: Record<string, unknown>
|
||||
}): ApiStreamUsageChunk {
|
||||
return {
|
||||
type: "usage",
|
||||
inputTokens: usage.inputTokens || 0,
|
||||
outputTokens: usage.outputTokens || 0,
|
||||
cacheReadTokens: usage.details?.cachedInputTokens,
|
||||
reasoningTokens: usage.details?.reasoningTokens,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Map OpenAI tool_choice to AI SDK toolChoice format.
|
||||
*/
|
||||
protected mapToolChoice(
|
||||
toolChoice: OpenAI.Chat.ChatCompletionCreateParams["tool_choice"],
|
||||
): "auto" | "none" | "required" | { type: "tool"; toolName: string } | undefined {
|
||||
if (!toolChoice) {
|
||||
return undefined
|
||||
}
|
||||
|
||||
// Handle string values
|
||||
if (typeof toolChoice === "string") {
|
||||
switch (toolChoice) {
|
||||
case "auto":
|
||||
return "auto"
|
||||
case "none":
|
||||
return "none"
|
||||
case "required":
|
||||
return "required"
|
||||
default:
|
||||
return "auto"
|
||||
}
|
||||
}
|
||||
|
||||
// Handle object values (OpenAI ChatCompletionNamedToolChoice format)
|
||||
if (typeof toolChoice === "object" && "type" in toolChoice) {
|
||||
if (toolChoice.type === "function" && "function" in toolChoice && toolChoice.function?.name) {
|
||||
return { type: "tool", toolName: toolChoice.function.name }
|
||||
}
|
||||
}
|
||||
|
||||
return undefined
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the max tokens parameter to include in the request.
|
||||
*/
|
||||
protected getMaxOutputTokens(): number | undefined {
|
||||
const modelInfo = this.config.modelInfo
|
||||
const maxTokens = this.config.modelMaxTokens || modelInfo.maxTokens
|
||||
|
||||
return maxTokens ?? undefined
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a message stream using the AI SDK.
|
||||
*/
|
||||
override async *createMessage(
|
||||
systemPrompt: string,
|
||||
messages: Anthropic.Messages.MessageParam[],
|
||||
metadata?: ApiHandlerCreateMessageMetadata,
|
||||
): ApiStream {
|
||||
const model = this.getModel()
|
||||
const languageModel = this.getLanguageModel()
|
||||
|
||||
// Convert messages to AI SDK format
|
||||
const aiSdkMessages = convertToAiSdkMessages(messages)
|
||||
|
||||
// Convert tools to OpenAI format first, then to AI SDK format
|
||||
const openAiTools = this.convertToolsForOpenAI(metadata?.tools)
|
||||
const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined
|
||||
|
||||
// Build the request options
|
||||
const requestOptions: Parameters<typeof streamText>[0] = {
|
||||
model: languageModel,
|
||||
system: systemPrompt,
|
||||
messages: aiSdkMessages,
|
||||
temperature: model.temperature ?? this.config.temperature ?? 0,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
tools: aiSdkTools,
|
||||
toolChoice: this.mapToolChoice(metadata?.tool_choice),
|
||||
}
|
||||
|
||||
// Use streamText for streaming responses
|
||||
const result = streamText(requestOptions)
|
||||
|
||||
// Process the full stream to get all events
|
||||
for await (const part of result.fullStream) {
|
||||
// Use the processAiSdkStreamPart utility to convert stream parts
|
||||
for (const chunk of processAiSdkStreamPart(part)) {
|
||||
yield chunk
|
||||
}
|
||||
}
|
||||
|
||||
// Yield usage metrics at the end
|
||||
const usage = await result.usage
|
||||
if (usage) {
|
||||
yield this.processUsageMetrics(usage)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Complete a prompt using the AI SDK generateText.
|
||||
*/
|
||||
async completePrompt(prompt: string): Promise<string> {
|
||||
const languageModel = this.getLanguageModel()
|
||||
|
||||
const { text } = await generateText({
|
||||
model: languageModel,
|
||||
prompt,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
temperature: this.config.temperature ?? 0,
|
||||
})
|
||||
|
||||
return text
|
||||
}
|
||||
}
|
||||
|
|
@ -532,6 +532,7 @@
|
|||
"zod": "3.25.76"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@ai-sdk/openai-compatible": "^1.0.0",
|
||||
"@openrouter/ai-sdk-provider": "^2.0.4",
|
||||
"@roo-code/build": "workspace:^",
|
||||
"@roo-code/config-eslint": "workspace:^",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue