Add native tool support for vercel ai gateway (#9697)

Co-authored-by: Roo Code <roomote@roocode.com>
This commit is contained in:
Matt Rubens 2025-12-01 14:52:05 -05:00 committed by GitHub
parent eb120543e1
commit b5acebccf5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 232 additions and 0 deletions

View file

@ -274,6 +274,217 @@ describe("VercelAiGatewayHandler", () => {
totalCost: 0.005,
})
})
describe("native tool calling", () => {
const testTools = [
{
type: "function" as const,
function: {
name: "test_tool",
description: "A test tool",
parameters: {
type: "object",
properties: {
arg1: { type: "string" },
},
required: ["arg1"],
},
},
},
]
beforeEach(() => {
mockCreate.mockImplementation(async () => ({
[Symbol.asyncIterator]: async function* () {
yield {
choices: [
{
delta: {},
index: 0,
},
],
}
},
}))
})
it("should include tools when provided", async () => {
const handler = new VercelAiGatewayHandler(mockOptions)
const messageGenerator = handler.createMessage("test prompt", [], {
taskId: "test-task-id",
tools: testTools,
toolProtocol: "native",
})
await messageGenerator.next()
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
tools: expect.arrayContaining([
expect.objectContaining({
type: "function",
function: expect.objectContaining({
name: "test_tool",
}),
}),
]),
}),
)
})
it("should include tool_choice when provided", async () => {
const handler = new VercelAiGatewayHandler(mockOptions)
const messageGenerator = handler.createMessage("test prompt", [], {
taskId: "test-task-id",
tools: testTools,
toolProtocol: "native",
tool_choice: "auto",
})
await messageGenerator.next()
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
tool_choice: "auto",
}),
)
})
it("should set parallel_tool_calls when toolProtocol is native", async () => {
const handler = new VercelAiGatewayHandler(mockOptions)
const messageGenerator = handler.createMessage("test prompt", [], {
taskId: "test-task-id",
tools: testTools,
toolProtocol: "native",
parallelToolCalls: true,
})
await messageGenerator.next()
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
parallel_tool_calls: true,
}),
)
})
it("should default parallel_tool_calls to false", async () => {
const handler = new VercelAiGatewayHandler(mockOptions)
const messageGenerator = handler.createMessage("test prompt", [], {
taskId: "test-task-id",
tools: testTools,
toolProtocol: "native",
})
await messageGenerator.next()
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
parallel_tool_calls: false,
}),
)
})
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)
const stream = handler.createMessage("test prompt", [], {
taskId: "test-task-id",
tools: testTools,
toolProtocol: "native",
})
const chunks = []
for await (const chunk of stream) {
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,
id: "call_123",
name: "test_tool",
arguments: '{"arg1":',
})
expect(toolCallChunks[1]).toEqual({
type: "tool_call_partial",
index: 0,
id: undefined,
name: undefined,
arguments: '"value"}',
})
})
it("should include stream_options with include_usage", async () => {
const handler = new VercelAiGatewayHandler(mockOptions)
const messageGenerator = handler.createMessage("test prompt", [], {
taskId: "test-task-id",
})
await messageGenerator.next()
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
stream_options: { include_usage: true },
}),
)
})
})
})
describe("completePrompt", () => {

View file

@ -176,6 +176,7 @@ describe("Vercel AI Gateway Fetchers", () => {
maxTokens: 8000,
contextWindow: 100000,
supportsImages: false,
supportsNativeTools: true,
supportsPromptCache: false,
inputPrice: 2500000,
outputPrice: 10000000,

View file

@ -108,6 +108,7 @@ export const parseVercelAiGatewayModel = ({ id, model }: { id: string; model: Ve
contextWindow: model.context_window,
supportsImages,
supportsPromptCache,
supportsNativeTools: true,
inputPrice: parseApiPrice(model.pricing?.input),
outputPrice: parseApiPrice(model.pricing?.output),
cacheWritesPrice,

View file

@ -60,6 +60,12 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp
: undefined,
max_completion_tokens: info.maxTokens,
stream: true,
stream_options: { include_usage: true },
...(metadata?.tools && { tools: this.convertToolsForOpenAI(metadata.tools) }),
...(metadata?.tool_choice && { tool_choice: metadata.tool_choice }),
...(metadata?.toolProtocol === "native" && {
parallel_tool_calls: metadata.parallelToolCalls ?? false,
}),
}
const completion = await this.client.chat.completions.create(body)
@ -73,6 +79,19 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp
}
}
// 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 {