mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
feat: migrate VercelAiGatewayHandler to AI SDK (#11353)
This commit is contained in:
Daniel
GitHub
parent
a4914c438c
commit
34a278e9a2
3 changed files with 319 additions and 431 deletions
|
|
@ -1,20 +1,28 @@
|
|||
// npx vitest run src/api/providers/__tests__/vercel-ai-gateway.spec.ts
|
||||
|
||||
// Mock vscode first to avoid import errors
|
||||
vitest.mock("vscode", () => ({}))
|
||||
// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
|
||||
const { mockStreamText, mockGenerateText, mockCreateGateway } = vi.hoisted(() => ({
|
||||
mockStreamText: vi.fn(),
|
||||
mockGenerateText: vi.fn(),
|
||||
mockCreateGateway: vi.fn(),
|
||||
}))
|
||||
|
||||
import { Anthropic } from "@anthropic-ai/sdk"
|
||||
import OpenAI from "openai"
|
||||
vi.mock("vscode", () => ({}))
|
||||
|
||||
import { VercelAiGatewayHandler } from "../vercel-ai-gateway"
|
||||
import { ApiHandlerOptions } from "../../../shared/api"
|
||||
import { vercelAiGatewayDefaultModelId, VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types"
|
||||
vi.mock("ai", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("ai")>()
|
||||
return {
|
||||
...actual,
|
||||
streamText: mockStreamText,
|
||||
generateText: mockGenerateText,
|
||||
createGateway: mockCreateGateway,
|
||||
}
|
||||
})
|
||||
|
||||
// Mock dependencies
|
||||
vitest.mock("openai")
|
||||
vitest.mock("delay", () => ({ default: vitest.fn(() => Promise.resolve()) }))
|
||||
vitest.mock("../fetchers/modelCache", () => ({
|
||||
getModels: vitest.fn().mockImplementation(() => {
|
||||
vi.mock("delay", () => ({ default: vi.fn(() => Promise.resolve()) }))
|
||||
|
||||
vi.mock("../fetchers/modelCache", () => ({
|
||||
getModels: vi.fn().mockImplementation(() => {
|
||||
return Promise.resolve({
|
||||
"anthropic/claude-sonnet-4": {
|
||||
maxTokens: 64000,
|
||||
|
|
@ -51,30 +59,22 @@ vitest.mock("../fetchers/modelCache", () => ({
|
|||
},
|
||||
})
|
||||
}),
|
||||
getModelsFromCache: vitest.fn().mockReturnValue(undefined),
|
||||
getModelsFromCache: vi.fn().mockReturnValue(undefined),
|
||||
}))
|
||||
|
||||
vitest.mock("../../transform/caching/vercel-ai-gateway", () => ({
|
||||
addCacheBreakpoints: vitest.fn(),
|
||||
import type { Anthropic } from "@anthropic-ai/sdk"
|
||||
|
||||
import { VercelAiGatewayHandler } from "../vercel-ai-gateway"
|
||||
import type { ApiHandlerOptions } from "../../../shared/api"
|
||||
import { vercelAiGatewayDefaultModelId, VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types"
|
||||
|
||||
// Set up the createGateway mock to return a function that creates mock language models
|
||||
const mockGatewayProvider = vi.fn((modelId: string) => ({
|
||||
modelId,
|
||||
provider: "gateway",
|
||||
}))
|
||||
|
||||
const mockCreate = vitest.fn()
|
||||
const mockConstructor = vitest.fn()
|
||||
|
||||
;(OpenAI as any).mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
}))
|
||||
;(OpenAI as any).mockImplementation = mockConstructor.mockReturnValue({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
})
|
||||
mockCreateGateway.mockReturnValue(mockGatewayProvider)
|
||||
|
||||
describe("VercelAiGatewayHandler", () => {
|
||||
const mockOptions: ApiHandlerOptions = {
|
||||
|
|
@ -83,19 +83,17 @@ describe("VercelAiGatewayHandler", () => {
|
|||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vitest.clearAllMocks()
|
||||
mockCreate.mockClear()
|
||||
mockConstructor.mockClear()
|
||||
vi.clearAllMocks()
|
||||
mockCreateGateway.mockReturnValue(mockGatewayProvider)
|
||||
})
|
||||
|
||||
it("initializes with correct options", () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
expect(handler).toBeInstanceOf(VercelAiGatewayHandler)
|
||||
|
||||
expect(OpenAI).toHaveBeenCalledWith({
|
||||
baseURL: "https://ai-gateway.vercel.sh/v1",
|
||||
expect(mockCreateGateway).toHaveBeenCalledWith({
|
||||
apiKey: mockOptions.vercelAiGatewayApiKey,
|
||||
defaultHeaders: expect.objectContaining({
|
||||
headers: expect.objectContaining({
|
||||
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline",
|
||||
"X-Title": "Roo Code",
|
||||
"User-Agent": expect.stringContaining("RooCode/"),
|
||||
|
|
@ -103,6 +101,11 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
})
|
||||
|
||||
it("reports as AI SDK provider", () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
expect(handler.isAiSdkProvider()).toBe(true)
|
||||
})
|
||||
|
||||
describe("fetchModel", () => {
|
||||
it("returns correct model info when options are provided", async () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
|
@ -130,41 +133,41 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
|
||||
describe("createMessage", () => {
|
||||
beforeEach(() => {
|
||||
mockCreate.mockImplementation(async () => ({
|
||||
[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,
|
||||
cache_creation_input_tokens: 2,
|
||||
prompt_tokens_details: {
|
||||
cached_tokens: 3,
|
||||
},
|
||||
cost: 0.005,
|
||||
},
|
||||
}
|
||||
},
|
||||
}))
|
||||
})
|
||||
function createMockStreamResult(options?: {
|
||||
usage?: { inputTokens: number; outputTokens: number; details?: Record<string, unknown> }
|
||||
providerMetadata?: Record<string, Record<string, unknown>>
|
||||
fullStream?: AsyncGenerator<any>
|
||||
}) {
|
||||
const defaultUsage = {
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
}
|
||||
|
||||
async function* defaultFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
return {
|
||||
fullStream: options?.fullStream ?? defaultFullStream(),
|
||||
usage: Promise.resolve(options?.usage ?? defaultUsage),
|
||||
providerMetadata: Promise.resolve(options?.providerMetadata ?? {}),
|
||||
}
|
||||
}
|
||||
|
||||
it("streams text content correctly", async () => {
|
||||
mockStreamText.mockReturnValue(
|
||||
createMockStreamResult({
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
providerMetadata: {
|
||||
gateway: {
|
||||
cache_creation_input_tokens: 2,
|
||||
cached_tokens: 3,
|
||||
cost: 0.005,
|
||||
},
|
||||
},
|
||||
}),
|
||||
)
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
const systemPrompt = "You are a helpful assistant."
|
||||
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
|
||||
|
|
@ -192,6 +195,8 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
it("uses correct temperature from options", async () => {
|
||||
const customTemp = 0.5
|
||||
mockStreamText.mockReturnValue(createMockStreamResult())
|
||||
|
||||
const handler = new VercelAiGatewayHandler({
|
||||
...mockOptions,
|
||||
modelTemperature: customTemp,
|
||||
|
|
@ -202,7 +207,7 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
await handler.createMessage(systemPrompt, messages).next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
temperature: customTemp,
|
||||
}),
|
||||
|
|
@ -210,6 +215,8 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
|
||||
it("uses default temperature when none provided", async () => {
|
||||
mockStreamText.mockReturnValue(createMockStreamResult())
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
const systemPrompt = "You are a helpful assistant."
|
||||
|
|
@ -217,29 +224,16 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
await handler.createMessage(systemPrompt, messages).next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
temperature: VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("adds cache breakpoints for supported models", async () => {
|
||||
const { addCacheBreakpoints } = await import("../../transform/caching/vercel-ai-gateway")
|
||||
const handler = new VercelAiGatewayHandler({
|
||||
...mockOptions,
|
||||
vercelAiGatewayModelId: "anthropic/claude-3.5-haiku",
|
||||
})
|
||||
it("sets correct maxOutputTokens", async () => {
|
||||
mockStreamText.mockReturnValue(createMockStreamResult())
|
||||
|
||||
const systemPrompt = "You are a helpful assistant."
|
||||
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
|
||||
|
||||
await handler.createMessage(systemPrompt, messages).next()
|
||||
|
||||
expect(addCacheBreakpoints).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("sets correct max_completion_tokens", async () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
const systemPrompt = "You are a helpful assistant."
|
||||
|
|
@ -247,14 +241,27 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
await handler.createMessage(systemPrompt, messages).next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
max_completion_tokens: 64000, // max tokens for sonnet 4
|
||||
maxOutputTokens: 64000,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("handles usage info correctly with all Vercel AI Gateway specific fields", async () => {
|
||||
mockStreamText.mockReturnValue(
|
||||
createMockStreamResult({
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
providerMetadata: {
|
||||
gateway: {
|
||||
cache_creation_input_tokens: 2,
|
||||
cached_tokens: 3,
|
||||
cost: 0.005,
|
||||
},
|
||||
},
|
||||
}),
|
||||
)
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
const systemPrompt = "You are a helpful assistant."
|
||||
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
|
||||
|
|
@ -294,22 +301,9 @@ describe("VercelAiGatewayHandler", () => {
|
|||
},
|
||||
]
|
||||
|
||||
beforeEach(() => {
|
||||
mockCreate.mockImplementation(async () => ({
|
||||
[Symbol.asyncIterator]: async function* () {
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: {},
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
}
|
||||
},
|
||||
}))
|
||||
})
|
||||
|
||||
it("should include tools when provided", async () => {
|
||||
mockStreamText.mockReturnValue(createMockStreamResult())
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
const messageGenerator = handler.createMessage("test prompt", [], {
|
||||
|
|
@ -318,21 +312,18 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
await messageGenerator.next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tools: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
type: "function",
|
||||
function: expect.objectContaining({
|
||||
name: "test_tool",
|
||||
}),
|
||||
}),
|
||||
]),
|
||||
tools: expect.objectContaining({
|
||||
test_tool: expect.any(Object),
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should include tool_choice when provided", async () => {
|
||||
it("should include toolChoice when provided", async () => {
|
||||
mockStreamText.mockReturnValue(createMockStreamResult())
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
const messageGenerator = handler.createMessage("test prompt", [], {
|
||||
|
|
@ -342,100 +333,42 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
await messageGenerator.next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tool_choice: "auto",
|
||||
toolChoice: "auto",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should set parallel_tool_calls when parallelToolCalls is enabled", async () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
it("should yield tool call events when streaming tool calls", async () => {
|
||||
async function* toolCallStream() {
|
||||
yield {
|
||||
type: "tool-input-start",
|
||||
id: "call_123",
|
||||
toolName: "test_tool",
|
||||
}
|
||||
yield {
|
||||
type: "tool-input-delta",
|
||||
id: "call_123",
|
||||
delta: '{"arg1":',
|
||||
}
|
||||
yield {
|
||||
type: "tool-input-delta",
|
||||
id: "call_123",
|
||||
delta: '"value"}',
|
||||
}
|
||||
yield {
|
||||
type: "tool-input-end",
|
||||
id: "call_123",
|
||||
}
|
||||
}
|
||||
|
||||
const messageGenerator = handler.createMessage("test prompt", [], {
|
||||
taskId: "test-task-id",
|
||||
tools: testTools,
|
||||
parallelToolCalls: true,
|
||||
})
|
||||
await messageGenerator.next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
parallel_tool_calls: true,
|
||||
mockStreamText.mockReturnValue(
|
||||
createMockStreamResult({
|
||||
fullStream: toolCallStream(),
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should include parallel_tool_calls: true by default", async () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
const messageGenerator = handler.createMessage("test prompt", [], {
|
||||
taskId: "test-task-id",
|
||||
tools: testTools,
|
||||
})
|
||||
await messageGenerator.next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tools: expect.any(Array),
|
||||
parallel_tool_calls: true,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should yield tool_call_partial chunks when streaming tool calls", async () => {
|
||||
mockCreate.mockImplementation(async () => ({
|
||||
[Symbol.asyncIterator]: async function* () {
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: {
|
||||
tool_calls: [
|
||||
{
|
||||
index: 0,
|
||||
id: "call_123",
|
||||
function: {
|
||||
name: "test_tool",
|
||||
arguments: '{"arg1":',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
}
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: {
|
||||
tool_calls: [
|
||||
{
|
||||
index: 0,
|
||||
function: {
|
||||
arguments: '"value"}',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
}
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: {},
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 5,
|
||||
},
|
||||
}
|
||||
},
|
||||
}))
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
|
|
@ -449,25 +382,38 @@ describe("VercelAiGatewayHandler", () => {
|
|||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
|
||||
expect(toolCallChunks).toHaveLength(2)
|
||||
expect(toolCallChunks[0]).toEqual({
|
||||
type: "tool_call_partial",
|
||||
index: 0,
|
||||
const toolStartChunks = chunks.filter((chunk) => chunk.type === "tool_call_start")
|
||||
expect(toolStartChunks).toHaveLength(1)
|
||||
expect(toolStartChunks[0]).toEqual({
|
||||
type: "tool_call_start",
|
||||
id: "call_123",
|
||||
name: "test_tool",
|
||||
arguments: '{"arg1":',
|
||||
})
|
||||
expect(toolCallChunks[1]).toEqual({
|
||||
type: "tool_call_partial",
|
||||
index: 0,
|
||||
id: undefined,
|
||||
name: undefined,
|
||||
arguments: '"value"}',
|
||||
|
||||
const toolDeltaChunks = chunks.filter((chunk) => chunk.type === "tool_call_delta")
|
||||
expect(toolDeltaChunks).toHaveLength(2)
|
||||
expect(toolDeltaChunks[0]).toEqual({
|
||||
type: "tool_call_delta",
|
||||
id: "call_123",
|
||||
delta: '{"arg1":',
|
||||
})
|
||||
expect(toolDeltaChunks[1]).toEqual({
|
||||
type: "tool_call_delta",
|
||||
id: "call_123",
|
||||
delta: '"value"}',
|
||||
})
|
||||
|
||||
const toolEndChunks = chunks.filter((chunk) => chunk.type === "tool_call_end")
|
||||
expect(toolEndChunks).toHaveLength(1)
|
||||
expect(toolEndChunks[0]).toEqual({
|
||||
type: "tool_call_end",
|
||||
id: "call_123",
|
||||
})
|
||||
})
|
||||
|
||||
it("should include stream_options with include_usage", async () => {
|
||||
it("should pass system prompt to streamText", async () => {
|
||||
mockStreamText.mockReturnValue(createMockStreamResult())
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
const messageGenerator = handler.createMessage("test prompt", [], {
|
||||
|
|
@ -475,9 +421,9 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
await messageGenerator.next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
stream_options: { include_usage: true },
|
||||
system: "test prompt",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
|
@ -485,43 +431,32 @@ describe("VercelAiGatewayHandler", () => {
|
|||
})
|
||||
|
||||
describe("completePrompt", () => {
|
||||
beforeEach(() => {
|
||||
mockCreate.mockImplementation(async () => ({
|
||||
choices: [
|
||||
{
|
||||
message: { role: "assistant", content: "Test completion response" },
|
||||
finish_reason: "stop",
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 8,
|
||||
completion_tokens: 4,
|
||||
total_tokens: 12,
|
||||
},
|
||||
}))
|
||||
})
|
||||
|
||||
it("completes prompt correctly", async () => {
|
||||
mockGenerateText.mockResolvedValue({
|
||||
text: "Test completion response",
|
||||
})
|
||||
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
const prompt = "Complete this: Hello"
|
||||
|
||||
const result = await handler.completePrompt(prompt)
|
||||
|
||||
expect(result).toBe("Test completion response")
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockGenerateText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
model: "anthropic/claude-sonnet-4",
|
||||
messages: [{ role: "user", content: prompt }],
|
||||
stream: false,
|
||||
prompt,
|
||||
temperature: VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE,
|
||||
max_completion_tokens: 64000,
|
||||
maxOutputTokens: 64000,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("uses custom temperature for completion", async () => {
|
||||
const customTemp = 0.8
|
||||
mockGenerateText.mockResolvedValue({
|
||||
text: "Test completion response",
|
||||
})
|
||||
|
||||
const handler = new VercelAiGatewayHandler({
|
||||
...mockOptions,
|
||||
modelTemperature: customTemp,
|
||||
|
|
@ -529,7 +464,7 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
await handler.completePrompt("Test prompt")
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockGenerateText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
temperature: customTemp,
|
||||
}),
|
||||
|
|
@ -540,27 +475,17 @@ describe("VercelAiGatewayHandler", () => {
|
|||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
const errorMessage = "API error"
|
||||
|
||||
mockCreate.mockImplementation(() => {
|
||||
throw new Error(errorMessage)
|
||||
})
|
||||
mockGenerateText.mockRejectedValue(new Error(errorMessage))
|
||||
|
||||
await expect(handler.completePrompt("Test")).rejects.toThrow(
|
||||
`Vercel AI Gateway completion error: ${errorMessage}`,
|
||||
)
|
||||
await expect(handler.completePrompt("Test")).rejects.toThrow("Vercel AI Gateway")
|
||||
})
|
||||
|
||||
it("returns empty string when no content in response", async () => {
|
||||
it("returns empty string when generateText returns empty text", async () => {
|
||||
const handler = new VercelAiGatewayHandler(mockOptions)
|
||||
|
||||
mockCreate.mockImplementation(async () => ({
|
||||
choices: [
|
||||
{
|
||||
message: { role: "assistant", content: null },
|
||||
finish_reason: "stop",
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
}))
|
||||
mockGenerateText.mockResolvedValue({
|
||||
text: "",
|
||||
})
|
||||
|
||||
const result = await handler.completePrompt("Test")
|
||||
expect(result).toBe("")
|
||||
|
|
@ -569,6 +494,10 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
describe("temperature support", () => {
|
||||
it("applies temperature for supported models", async () => {
|
||||
mockGenerateText.mockResolvedValue({
|
||||
text: "Test response",
|
||||
})
|
||||
|
||||
const handler = new VercelAiGatewayHandler({
|
||||
...mockOptions,
|
||||
vercelAiGatewayModelId: "anthropic/claude-sonnet-4",
|
||||
|
|
@ -577,7 +506,7 @@ describe("VercelAiGatewayHandler", () => {
|
|||
|
||||
await handler.completePrompt("Test")
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockGenerateText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
temperature: 0.9,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -1,87 +0,0 @@
|
|||
import OpenAI from "openai"
|
||||
|
||||
import { type ModelInfo, type ModelRecord } from "@roo-code/types"
|
||||
|
||||
import { ApiHandlerOptions, RouterName } from "../../shared/api"
|
||||
|
||||
import { BaseProvider } from "./base-provider"
|
||||
import { getModels, getModelsFromCache } from "./fetchers/modelCache"
|
||||
|
||||
import { DEFAULT_HEADERS } from "./constants"
|
||||
|
||||
type RouterProviderOptions = {
|
||||
name: RouterName
|
||||
baseURL: string
|
||||
apiKey?: string
|
||||
modelId?: string
|
||||
defaultModelId: string
|
||||
defaultModelInfo: ModelInfo
|
||||
options: ApiHandlerOptions
|
||||
}
|
||||
|
||||
export abstract class RouterProvider extends BaseProvider {
|
||||
protected readonly options: ApiHandlerOptions
|
||||
protected readonly name: RouterName
|
||||
protected models: ModelRecord = {}
|
||||
protected readonly modelId?: string
|
||||
protected readonly defaultModelId: string
|
||||
protected readonly defaultModelInfo: ModelInfo
|
||||
protected readonly client: OpenAI
|
||||
|
||||
constructor({
|
||||
options,
|
||||
name,
|
||||
baseURL,
|
||||
apiKey = "not-provided",
|
||||
modelId,
|
||||
defaultModelId,
|
||||
defaultModelInfo,
|
||||
}: RouterProviderOptions) {
|
||||
super()
|
||||
|
||||
this.options = options
|
||||
this.name = name
|
||||
this.modelId = modelId
|
||||
this.defaultModelId = defaultModelId
|
||||
this.defaultModelInfo = defaultModelInfo
|
||||
|
||||
this.client = new OpenAI({
|
||||
baseURL,
|
||||
apiKey,
|
||||
defaultHeaders: {
|
||||
...DEFAULT_HEADERS,
|
||||
...(options.openAiHeaders || {}),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
public async fetchModel() {
|
||||
this.models = await getModels({ provider: this.name, apiKey: this.client.apiKey, baseUrl: this.client.baseURL })
|
||||
return this.getModel()
|
||||
}
|
||||
|
||||
override getModel(): { id: string; info: ModelInfo } {
|
||||
const id = this.modelId ?? this.defaultModelId
|
||||
|
||||
// First check instance models (populated by fetchModel)
|
||||
if (this.models[id]) {
|
||||
return { id, info: this.models[id] }
|
||||
}
|
||||
|
||||
// Fall back to global cache (synchronous disk/memory cache)
|
||||
// This ensures models are available before fetchModel() is called
|
||||
const cachedModels = getModelsFromCache(this.name)
|
||||
if (cachedModels?.[id]) {
|
||||
// Also populate instance models for future calls
|
||||
this.models = cachedModels
|
||||
return { id, info: cachedModels[id] }
|
||||
}
|
||||
|
||||
// Last resort: return default model
|
||||
return { id: this.defaultModelId, info: this.defaultModelInfo }
|
||||
}
|
||||
|
||||
protected supportsTemperature(modelId: string): boolean {
|
||||
return !modelId.startsWith("openai/o3-mini")
|
||||
}
|
||||
}
|
||||
|
|
@ -1,132 +1,178 @@
|
|||
import { Anthropic } from "@anthropic-ai/sdk"
|
||||
import OpenAI from "openai"
|
||||
import { createGateway, streamText, generateText, ToolSet } from "ai"
|
||||
|
||||
import {
|
||||
vercelAiGatewayDefaultModelId,
|
||||
vercelAiGatewayDefaultModelInfo,
|
||||
VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE,
|
||||
VERCEL_AI_GATEWAY_PROMPT_CACHING_MODELS,
|
||||
type ModelInfo,
|
||||
type ModelRecord,
|
||||
} from "@roo-code/types"
|
||||
|
||||
import { ApiHandlerOptions } from "../../shared/api"
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
||||
import { ApiStream } from "../transform/stream"
|
||||
import { convertToOpenAiMessages } from "../transform/openai-format"
|
||||
import { addCacheBreakpoints } from "../transform/caching/vercel-ai-gateway"
|
||||
import {
|
||||
convertToAiSdkMessages,
|
||||
convertToolsForAiSdk,
|
||||
processAiSdkStreamPart,
|
||||
mapToolChoice,
|
||||
handleAiSdkError,
|
||||
} from "../transform/ai-sdk"
|
||||
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
|
||||
|
||||
import { DEFAULT_HEADERS } from "./constants"
|
||||
import { BaseProvider } from "./base-provider"
|
||||
import { getModels, getModelsFromCache } from "./fetchers/modelCache"
|
||||
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
|
||||
import { RouterProvider } from "./router-provider"
|
||||
|
||||
// Extend OpenAI's CompletionUsage to include Vercel AI Gateway specific fields
|
||||
interface VercelAiGatewayUsage extends OpenAI.CompletionUsage {
|
||||
cache_creation_input_tokens?: number
|
||||
cost?: number
|
||||
}
|
||||
/**
|
||||
* Vercel AI Gateway provider using the built-in AI SDK gateway support.
|
||||
* Uses `createGateway` from the `ai` package to communicate with the
|
||||
* Vercel AI Gateway v3 API at https://ai-gateway.vercel.sh/v3/ai.
|
||||
*/
|
||||
export class VercelAiGatewayHandler extends BaseProvider implements SingleCompletionHandler {
|
||||
protected options: ApiHandlerOptions
|
||||
protected provider: ReturnType<typeof createGateway>
|
||||
private readonly name = "vercel-ai-gateway" as const
|
||||
private models: ModelRecord = {}
|
||||
|
||||
export class VercelAiGatewayHandler extends RouterProvider implements SingleCompletionHandler {
|
||||
constructor(options: ApiHandlerOptions) {
|
||||
super({
|
||||
options,
|
||||
name: "vercel-ai-gateway",
|
||||
baseURL: "https://ai-gateway.vercel.sh/v1",
|
||||
apiKey: options.vercelAiGatewayApiKey,
|
||||
modelId: options.vercelAiGatewayModelId,
|
||||
defaultModelId: vercelAiGatewayDefaultModelId,
|
||||
defaultModelInfo: vercelAiGatewayDefaultModelInfo,
|
||||
super()
|
||||
this.options = options
|
||||
|
||||
this.provider = createGateway({
|
||||
apiKey: options.vercelAiGatewayApiKey ?? "not-provided",
|
||||
headers: DEFAULT_HEADERS,
|
||||
})
|
||||
}
|
||||
|
||||
override getModel(): { id: string; info: ModelInfo } {
|
||||
const id = this.options.vercelAiGatewayModelId ?? vercelAiGatewayDefaultModelId
|
||||
|
||||
if (this.models[id]) {
|
||||
return { id, info: this.models[id] }
|
||||
}
|
||||
|
||||
const cachedModels = getModelsFromCache(this.name)
|
||||
if (cachedModels?.[id]) {
|
||||
this.models = cachedModels
|
||||
return { id, info: cachedModels[id] }
|
||||
}
|
||||
|
||||
return { id: vercelAiGatewayDefaultModelId, info: vercelAiGatewayDefaultModelInfo }
|
||||
}
|
||||
|
||||
public async fetchModel() {
|
||||
this.models = await getModels({
|
||||
provider: this.name,
|
||||
apiKey: this.options.vercelAiGatewayApiKey ?? "not-provided",
|
||||
})
|
||||
return this.getModel()
|
||||
}
|
||||
|
||||
protected getLanguageModel(modelId?: string) {
|
||||
const id = modelId ?? this.getModel().id
|
||||
return this.provider(id)
|
||||
}
|
||||
|
||||
protected supportsTemperature(modelId: string): boolean {
|
||||
return !modelId.startsWith("openai/o3-mini")
|
||||
}
|
||||
|
||||
protected processUsageMetrics(
|
||||
usage: {
|
||||
inputTokens?: number
|
||||
outputTokens?: number
|
||||
details?: {
|
||||
cachedInputTokens?: number
|
||||
reasoningTokens?: number
|
||||
}
|
||||
},
|
||||
providerMetadata?: Record<string, Record<string, unknown>>,
|
||||
): ApiStreamUsageChunk {
|
||||
const gatewayMeta = providerMetadata?.gateway as Record<string, unknown> | undefined
|
||||
|
||||
const cacheWriteTokens = (gatewayMeta?.cache_creation_input_tokens as number) ?? undefined
|
||||
const cacheReadTokens = usage.details?.cachedInputTokens ?? (gatewayMeta?.cached_tokens as number) ?? undefined
|
||||
const totalCost = (gatewayMeta?.cost as number) ?? 0
|
||||
|
||||
return {
|
||||
type: "usage",
|
||||
inputTokens: usage.inputTokens || 0,
|
||||
outputTokens: usage.outputTokens || 0,
|
||||
cacheWriteTokens,
|
||||
cacheReadTokens,
|
||||
totalCost,
|
||||
}
|
||||
}
|
||||
|
||||
override async *createMessage(
|
||||
systemPrompt: string,
|
||||
messages: Anthropic.Messages.MessageParam[],
|
||||
metadata?: ApiHandlerCreateMessageMetadata,
|
||||
): ApiStream {
|
||||
const { id: modelId, info } = await this.fetchModel()
|
||||
const languageModel = this.getLanguageModel(modelId)
|
||||
|
||||
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
|
||||
{ role: "system", content: systemPrompt },
|
||||
...convertToOpenAiMessages(messages),
|
||||
]
|
||||
const aiSdkMessages = convertToAiSdkMessages(messages)
|
||||
|
||||
if (VERCEL_AI_GATEWAY_PROMPT_CACHING_MODELS.has(modelId) && info.supportsPromptCache) {
|
||||
addCacheBreakpoints(systemPrompt, openAiMessages)
|
||||
}
|
||||
const openAiTools = this.convertToolsForOpenAI(metadata?.tools)
|
||||
const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined
|
||||
|
||||
const body: OpenAI.Chat.ChatCompletionCreateParams = {
|
||||
model: modelId,
|
||||
messages: openAiMessages,
|
||||
temperature: this.supportsTemperature(modelId)
|
||||
? (this.options.modelTemperature ?? VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE)
|
||||
: undefined,
|
||||
max_completion_tokens: info.maxTokens,
|
||||
stream: true,
|
||||
stream_options: { include_usage: true },
|
||||
tools: this.convertToolsForOpenAI(metadata?.tools),
|
||||
tool_choice: metadata?.tool_choice,
|
||||
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
|
||||
}
|
||||
const temperature = this.supportsTemperature(modelId)
|
||||
? (this.options.modelTemperature ?? VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE)
|
||||
: undefined
|
||||
|
||||
const completion = await this.client.chat.completions.create(body)
|
||||
const result = streamText({
|
||||
model: languageModel,
|
||||
system: systemPrompt,
|
||||
messages: aiSdkMessages,
|
||||
temperature,
|
||||
maxOutputTokens: info.maxTokens ?? undefined,
|
||||
tools: aiSdkTools,
|
||||
toolChoice: mapToolChoice(metadata?.tool_choice),
|
||||
})
|
||||
|
||||
for await (const chunk of completion) {
|
||||
const delta = chunk.choices[0]?.delta
|
||||
if (delta?.content) {
|
||||
yield {
|
||||
type: "text",
|
||||
text: delta.content,
|
||||
try {
|
||||
for await (const part of result.fullStream) {
|
||||
for (const chunk of processAiSdkStreamPart(part)) {
|
||||
yield chunk
|
||||
}
|
||||
}
|
||||
|
||||
// Emit raw tool call chunks - NativeToolCallParser handles state management
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (chunk.usage) {
|
||||
const usage = chunk.usage as VercelAiGatewayUsage
|
||||
yield {
|
||||
type: "usage",
|
||||
inputTokens: usage.prompt_tokens || 0,
|
||||
outputTokens: usage.completion_tokens || 0,
|
||||
cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
|
||||
cacheReadTokens: usage.prompt_tokens_details?.cached_tokens || undefined,
|
||||
totalCost: usage.cost ?? 0,
|
||||
}
|
||||
const usage = await result.usage
|
||||
const providerMetadata = await result.providerMetadata
|
||||
if (usage) {
|
||||
yield this.processUsageMetrics(usage, providerMetadata as any)
|
||||
}
|
||||
} catch (error) {
|
||||
throw handleAiSdkError(error, "Vercel AI Gateway")
|
||||
}
|
||||
}
|
||||
|
||||
async completePrompt(prompt: string): Promise<string> {
|
||||
const { id: modelId, info } = await this.fetchModel()
|
||||
const languageModel = this.getLanguageModel(modelId)
|
||||
|
||||
const temperature = this.supportsTemperature(modelId)
|
||||
? (this.options.modelTemperature ?? VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE)
|
||||
: undefined
|
||||
|
||||
try {
|
||||
const requestOptions: OpenAI.Chat.ChatCompletionCreateParams = {
|
||||
model: modelId,
|
||||
messages: [{ role: "user", content: prompt }],
|
||||
stream: false,
|
||||
}
|
||||
const { text } = await generateText({
|
||||
model: languageModel,
|
||||
prompt,
|
||||
maxOutputTokens: info.maxTokens ?? undefined,
|
||||
temperature,
|
||||
})
|
||||
|
||||
if (this.supportsTemperature(modelId)) {
|
||||
requestOptions.temperature = this.options.modelTemperature ?? VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE
|
||||
}
|
||||
|
||||
requestOptions.max_completion_tokens = info.maxTokens
|
||||
|
||||
const response = await this.client.chat.completions.create(requestOptions)
|
||||
return response.choices[0]?.message.content || ""
|
||||
return text
|
||||
} catch (error) {
|
||||
if (error instanceof Error) {
|
||||
throw new Error(`Vercel AI Gateway completion error: ${error.message}`)
|
||||
}
|
||||
throw error
|
||||
throw handleAiSdkError(error, "Vercel AI Gateway")
|
||||
}
|
||||
}
|
||||
|
||||
override isAiSdkProvider(): boolean {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue