scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support tool calling in native ollama provider (#9696)

Browse source 
Co-authored-by: Roo Code <roomote@roocode.com>
This commit is contained in:
Matt Rubens 2025-12-01 15:28:01 -05:00 • committed by GitHub
parent faa6c40ac1
commit be7659461a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 336 additions and 435 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
packages/types/src/providers/ollama.ts
View file

@ -8,6 +8,7 @@ export const ollamaDefaultModelInfo: ModelInfo = {
contextWindow: 200_000,
supportsImages: true,
supportsPromptCache: true,
supportsNativeTools: true,
inputPrice: 0,
outputPrice: 0,
cacheWritesPrice: 0,

278
src/api/providers/__tests__/native-ollama.spec.ts
View file

@ -2,6 +2,7 @@
import { NativeOllamaHandler } from "../native-ollama"
import { ApiHandlerOptions } from "../../../shared/api"
import { getOllamaModels } from "../fetchers/ollama"
// Mock the ollama package
const mockChat = vitest.fn()
@ -16,22 +17,27 @@ vitest.mock("ollama", () => {
// Mock the getOllamaModels function
vitest.mock("../fetchers/ollama", () => ({
getOllamaModels: vitest.fn().mockResolvedValue({
llama2: {
contextWindow: 4096,
maxTokens: 4096,
supportsImages: false,
supportsPromptCache: false,
},
}),
getOllamaModels: vitest.fn(),
}))
const mockGetOllamaModels = vitest.mocked(getOllamaModels)
describe("NativeOllamaHandler", () => {
let handler: NativeOllamaHandler
beforeEach(() => {
vitest.clearAllMocks()
// Default mock for getOllamaModels
mockGetOllamaModels.mockResolvedValue({
llama2: {
contextWindow: 4096,
maxTokens: 4096,
supportsImages: false,
supportsPromptCache: false,
},
})
const options: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
@ -257,4 +263,260 @@ describe("NativeOllamaHandler", () => {
expect(model.info).toBeDefined()
})
})
describe("tool calling", () => {
it("should include tools when model supports native tools", async () => {
// Mock model with native tool support
mockGetOllamaModels.mockResolvedValue({
"llama3.2": {
contextWindow: 128000,
maxTokens: 4096,
supportsImages: true,
supportsPromptCache: false,
supportsNativeTools: true,
},
})
const options: ApiHandlerOptions = {
apiModelId: "llama3.2",
ollamaModelId: "llama3.2",
ollamaBaseUrl: "http://localhost:11434",
}
handler = new NativeOllamaHandler(options)
// Mock the chat response
mockChat.mockImplementation(async function* () {
yield { message: { content: "I will use the tool" } }
})
const tools = [
{
type: "function" as const,
function: {
name: "get_weather",
description: "Get the weather for a location",
parameters: {
type: "object",
properties: {
location: { type: "string", description: "The city name" },
},
required: ["location"],
},
},
},
]
const stream = handler.createMessage(
"System",
[{ role: "user" as const, content: "What's the weather?" }],
{ taskId: "test", tools },
)
// Consume the stream
for await (const _ of stream) {
// consume stream
}
// Verify tools were passed to the API
expect(mockChat).toHaveBeenCalledWith(
expect.objectContaining({
tools: [
{
type: "function",
function: {
name: "get_weather",
description: "Get the weather for a location",
parameters: {
type: "object",
properties: {
location: { type: "string", description: "The city name" },
},
required: ["location"],
},
},
},
],
}),
)
})
it("should not include tools when model does not support native tools", async () => {
// Mock model without native tool support
mockGetOllamaModels.mockResolvedValue({
llama2: {
contextWindow: 4096,
maxTokens: 4096,
supportsImages: false,
supportsPromptCache: false,
supportsNativeTools: false,
},
})
// Mock the chat response
mockChat.mockImplementation(async function* () {
yield { message: { content: "Response without tools" } }
})
const tools = [
{
type: "function" as const,
function: {
name: "get_weather",
description: "Get the weather",
parameters: { type: "object", properties: {} },
},
},
]
const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }], {
taskId: "test",
tools,
})
// Consume the stream
for await (const _ of stream) {
// consume stream
}
// Verify tools were NOT passed
expect(mockChat).toHaveBeenCalledWith(
expect.not.objectContaining({
tools: expect.anything(),
}),
)
})
it("should not include tools when toolProtocol is xml", async () => {
// Mock model with native tool support
mockGetOllamaModels.mockResolvedValue({
"llama3.2": {
contextWindow: 128000,
maxTokens: 4096,
supportsImages: true,
supportsPromptCache: false,
supportsNativeTools: true,
},
})
const options: ApiHandlerOptions = {
apiModelId: "llama3.2",
ollamaModelId: "llama3.2",
ollamaBaseUrl: "http://localhost:11434",
}
handler = new NativeOllamaHandler(options)
// Mock the chat response
mockChat.mockImplementation(async function* () {
yield { message: { content: "Response" } }
})
const tools = [
{
type: "function" as const,
function: {
name: "get_weather",
description: "Get the weather",
parameters: { type: "object", properties: {} },
},
},
]
const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }], {
taskId: "test",
tools,
toolProtocol: "xml",
})
// Consume the stream
for await (const _ of stream) {
// consume stream
}
// Verify tools were NOT passed (XML protocol forces XML format)
expect(mockChat).toHaveBeenCalledWith(
expect.not.objectContaining({
tools: expect.anything(),
}),
)
})
it("should yield tool_call_partial when model returns tool calls", async () => {
// Mock model with native tool support
mockGetOllamaModels.mockResolvedValue({
"llama3.2": {
contextWindow: 128000,
maxTokens: 4096,
supportsImages: true,
supportsPromptCache: false,
supportsNativeTools: true,
},
})
const options: ApiHandlerOptions = {
apiModelId: "llama3.2",
ollamaModelId: "llama3.2",
ollamaBaseUrl: "http://localhost:11434",
}
handler = new NativeOllamaHandler(options)
// Mock the chat response with tool calls
mockChat.mockImplementation(async function* () {
yield {
message: {
content: "",
tool_calls: [
{
function: {
name: "get_weather",
arguments: { location: "San Francisco" },
},
},
],
},
}
})
const tools = [
{
type: "function" as const,
function: {
name: "get_weather",
description: "Get the weather for a location",
parameters: {
type: "object",
properties: {
location: { type: "string" },
},
required: ["location"],
},
},
},
]
const stream = handler.createMessage(
"System",
[{ role: "user" as const, content: "What's the weather in SF?" }],
{ taskId: "test", tools },
)
const results = []
for await (const chunk of stream) {
results.push(chunk)
}
// Should yield a tool_call_partial chunk
const toolCallChunk = results.find((r) => r.type === "tool_call_partial")
expect(toolCallChunk).toBeDefined()
expect(toolCallChunk).toEqual({
type: "tool_call_partial",
index: 0,
id: "ollama-tool-0",
name: "get_weather",
arguments: JSON.stringify({ location: "San Francisco" }),
})
})
})
})

108
src/api/providers/__tests__/ollama-timeout.spec.ts
View file

@ -1,108 +0,0 @@
// npx vitest run api/providers/__tests__/ollama-timeout.spec.ts
import { OllamaHandler } from "../ollama"
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("OllamaHandler timeout configuration", () => {
beforeEach(() => {
vitest.clearAllMocks()
})
it("should use default timeout of 600 seconds when no configuration is set", () => {
;(getApiRequestTimeout as any).mockReturnValue(600000)
const options: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
ollamaBaseUrl: "http://localhost:11434",
}
new OllamaHandler(options)
expect(getApiRequestTimeout).toHaveBeenCalled()
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
expect.objectContaining({
baseURL: "http://localhost:11434/v1",
apiKey: "ollama",
timeout: 600000, // 600 seconds in milliseconds
}),
)
})
it("should use custom timeout when configuration is set", () => {
;(getApiRequestTimeout as any).mockReturnValue(3600000) // 1 hour
const options: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
}
new OllamaHandler(options)
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
expect.objectContaining({
timeout: 3600000, // 3600 seconds in milliseconds
}),
)
})
it("should handle zero timeout (no timeout)", () => {
;(getApiRequestTimeout as any).mockReturnValue(0)
const options: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
ollamaBaseUrl: "http://localhost:11434",
}
new OllamaHandler(options)
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
expect.objectContaining({
timeout: 0, // No timeout
}),
)
})
it("should use default base URL when not provided", () => {
;(getApiRequestTimeout as any).mockReturnValue(600000)
const options: ApiHandlerOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
}
new OllamaHandler(options)
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
expect.objectContaining({
baseURL: "http://localhost:11434/v1",
}),
)
})
})

178
src/api/providers/__tests__/ollama.spec.ts
View file

@ -1,178 +0,0 @@
// npx vitest run api/providers/__tests__/ollama.spec.ts
import { Anthropic } from "@anthropic-ai/sdk"
import { OllamaHandler } from "../ollama"
import { ApiHandlerOptions } from "../../../shared/api"
const mockCreate = vitest.fn()
vitest.mock("openai", () => {
return {
__esModule: true,
default: vitest.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,
},
}
}
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,
},
}
},
}
}),
},
},
})),
}
})
describe("OllamaHandler", () => {
let handler: OllamaHandler
let mockOptions: ApiHandlerOptions
beforeEach(() => {
mockOptions = {
apiModelId: "llama2",
ollamaModelId: "llama2",
ollamaBaseUrl: "http://localhost:11434/v1",
}
handler = new OllamaHandler(mockOptions)
mockCreate.mockClear()
})
describe("constructor", () => {
it("should initialize with provided options", () => {
expect(handler).toBeInstanceOf(OllamaHandler)
expect(handler.getModel().id).toBe(mockOptions.ollamaModelId)
})
it("should use default base URL if not provided", () => {
const handlerWithoutUrl = new OllamaHandler({
apiModelId: "llama2",
ollamaModelId: "llama2",
})
expect(handlerWithoutUrl).toBeInstanceOf(OllamaHandler)
})
it("should use API key when provided", () => {
const handlerWithApiKey = new OllamaHandler({
apiModelId: "llama2",
ollamaModelId: "llama2",
ollamaBaseUrl: "https://ollama.com",
ollamaApiKey: "test-api-key",
})
expect(handlerWithApiKey).toBeInstanceOf(OllamaHandler)
// The API key will be used in the Authorization header
})
})
describe("createMessage", () => {
const systemPrompt = "You are a helpful assistant."
const messages: Anthropic.Messages.MessageParam[] = [
{
role: "user",
content: "Hello!",
},
]
it("should handle streaming responses", async () => {
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
expect(chunks.length).toBeGreaterThan(0)
const textChunks = chunks.filter((chunk) => chunk.type === "text")
expect(textChunks).toHaveLength(1)
expect(textChunks[0].text).toBe("Test response")
})
it("should handle API errors", async () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))
const stream = handler.createMessage(systemPrompt, messages)
await expect(async () => {
for await (const _chunk of stream) {
// Should not reach here
}
}).rejects.toThrow("API Error")
})
})
describe("completePrompt", () => {
it("should complete prompt successfully", async () => {
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("Test response")
expect(mockCreate).toHaveBeenCalledWith({
model: mockOptions.ollamaModelId,
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("Ollama completion error: API Error")
})
it("should handle empty response", async () => {
mockCreate.mockResolvedValueOnce({
choices: [{ message: { content: "" } }],
})
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("")
})
})
describe("getModel", () => {
it("should return model info", () => {
const modelInfo = handler.getModel()
expect(modelInfo.id).toBe(mockOptions.ollamaModelId)
expect(modelInfo.info).toBeDefined()
expect(modelInfo.info.maxTokens).toBe(-1)
expect(modelInfo.info.contextWindow).toBe(128_000)
})
})
})

2
src/api/providers/fetchers/__tests__/ollama.test.ts
View file

@ -22,6 +22,7 @@ describe("Ollama Fetcher", () => {
contextWindow: 40960,
supportsImages: false,
supportsPromptCache: true,
supportsNativeTools: true,
inputPrice: 0,
outputPrice: 0,
cacheWritesPrice: 0,
@ -46,6 +47,7 @@ describe("Ollama Fetcher", () => {
contextWindow: 40960,
supportsImages: false,
supportsPromptCache: true,
supportsNativeTools: true,
inputPrice: 0,
outputPrice: 0,
cacheWritesPrice: 0,

1
src/api/providers/index.ts
View file

@ -17,7 +17,6 @@ export { IOIntelligenceHandler } from "./io-intelligence"
export { LiteLLMHandler } from "./lite-llm"
export { LmStudioHandler } from "./lm-studio"
export { MistralHandler } from "./mistral"
export { OllamaHandler } from "./ollama"
export { OpenAiNativeHandler } from "./openai-native"
export { OpenAiHandler } from "./openai"
export { OpenRouterHandler } from "./openrouter"

66
src/api/providers/native-ollama.ts
View file

@ -1,5 +1,6 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { Message, Ollama, type Config as OllamaOptions } from "ollama"
import OpenAI from "openai"
import { Message, Ollama, Tool as OllamaTool, type Config as OllamaOptions } from "ollama"
import { ModelInfo, openAiModelInfoSaneDefaults, DEEP_SEEK_DEFAULT_TEMPERATURE } from "@roo-code/types"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
@ -93,7 +94,7 @@ function convertToOllamaMessages(anthropicMessages: Anthropic.Messages.MessagePa
})
}
} else if (anthropicMessage.role === "assistant") {
const { nonToolMessages } = anthropicMessage.content.reduce<{
const { nonToolMessages, toolMessages } = anthropicMessage.content.reduce<{
nonToolMessages: (Anthropic.TextBlockParam | Anthropic.ImageBlockParam)[]
toolMessages: Anthropic.ToolUseBlockParam[]
}>(
@ -121,9 +122,21 @@ function convertToOllamaMessages(anthropicMessages: Anthropic.Messages.MessagePa
.join("\n")
}
// Convert tool_use blocks to Ollama tool_calls format
const toolCalls =
toolMessages.length > 0
? toolMessages.map((tool) => ({
function: {
name: tool.name,
arguments: tool.input as Record<string, unknown>,
},
}))
: undefined
ollamaMessages.push({
role: "assistant",
content,
tool_calls: toolCalls,
})
}
}
@ -165,6 +178,28 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
return this.client
}
/**
* Converts OpenAI-format tools to Ollama's native tool format.
* This allows NativeOllamaHandler to use the same tool definitions
* that are passed to OpenAI-compatible providers.
*/
private convertToolsToOllama(tools: OpenAI.Chat.ChatCompletionTool[] | undefined): OllamaTool[] | undefined {
if (!tools || tools.length === 0) {
return undefined
}
return tools
.filter((tool): tool is OpenAI.Chat.ChatCompletionTool & { type: "function" } => tool.type === "function")
.map((tool) => ({
type: tool.type,
function: {
name: tool.function.name,
description: tool.function.description,
parameters: tool.function.parameters as OllamaTool["function"]["parameters"],
},
}))
}
override async *createMessage(
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
@ -188,6 +223,11 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
}) as const,
)
// Check if we should use native tool calling
const supportsNativeTools = modelInfo.supportsNativeTools ?? false
const useNativeTools =
supportsNativeTools && metadata?.tools && metadata.tools.length > 0 && metadata?.toolProtocol !== "xml"
try {
// Build options object conditionally
const chatOptions: OllamaChatOptions = {
@ -205,20 +245,40 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
messages: ollamaMessages,
stream: true,
options: chatOptions,
// Native tool calling support
...(useNativeTools && { tools: this.convertToolsToOllama(metadata.tools) }),
})
let totalInputTokens = 0
let totalOutputTokens = 0
// Track tool calls across chunks (Ollama may send complete tool_calls in final chunk)
let toolCallIndex = 0
try {
for await (const chunk of stream) {
if (typeof chunk.message.content === "string") {
if (typeof chunk.message.content === "string" && chunk.message.content.length > 0) {
// Process content through matcher for reasoning detection
for (const matcherChunk of matcher.update(chunk.message.content)) {
yield matcherChunk
}
}
// Handle tool calls - emit partial chunks for NativeToolCallParser compatibility
if (chunk.message.tool_calls && chunk.message.tool_calls.length > 0) {
for (const toolCall of chunk.message.tool_calls) {
// Generate a unique ID for this tool call
const toolCallId = `ollama-tool-${toolCallIndex}`
yield {
type: "tool_call_partial",
index: toolCallIndex,
id: toolCallId,
name: toolCall.function.name,
arguments: JSON.stringify(toolCall.function.arguments),
}
toolCallIndex++
}
}
// Handle token usage if available
if (chunk.eval_count !== undefined || chunk.prompt_eval_count !== undefined) {
if (chunk.prompt_eval_count) {

137
src/api/providers/ollama.ts
View file

@ -1,137 +0,0 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { type ModelInfo, openAiModelInfoSaneDefaults, DEEP_SEEK_DEFAULT_TEMPERATURE } from "@roo-code/types"
import type { ApiHandlerOptions } from "../../shared/api"
import { XmlMatcher } from "../../utils/xml-matcher"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { convertToR1Format } from "../transform/r1-format"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
import { getApiRequestTimeout } from "./utils/timeout-config"
import { handleOpenAIError } from "./utils/openai-error-handler"
type CompletionUsage = OpenAI.Chat.Completions.ChatCompletionChunk["usage"]
export class OllamaHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
private readonly providerName = "Ollama"
constructor(options: ApiHandlerOptions) {
super()
this.options = options
// Use the API key if provided (for Ollama cloud or authenticated instances)
// Otherwise use "ollama" as a placeholder for local instances
const apiKey = this.options.ollamaApiKey || "ollama"
const headers: Record<string, string> = {}
if (this.options.ollamaApiKey) {
headers["Authorization"] = `Bearer ${this.options.ollamaApiKey}`
}
this.client = new OpenAI({
baseURL: (this.options.ollamaBaseUrl || "http://localhost:11434") + "/v1",
apiKey: apiKey,
timeout: getApiRequestTimeout(),
defaultHeaders: headers,
})
}
override async *createMessage(
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
metadata?: ApiHandlerCreateMessageMetadata,
): ApiStream {
const modelId = this.getModel().id
const useR1Format = modelId.toLowerCase().includes("deepseek-r1")
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...(useR1Format ? convertToR1Format(messages) : convertToOpenAiMessages(messages)),
]
let stream
try {
stream = await this.client.chat.completions.create({
model: this.getModel().id,
messages: openAiMessages,
temperature: this.options.modelTemperature ?? 0,
stream: true,
stream_options: { include_usage: true },
})
} catch (error) {
throw handleOpenAIError(error, this.providerName)
}
const matcher = new XmlMatcher(
"think",
(chunk) =>
({
type: chunk.matched ? "reasoning" : "text",
text: chunk.data,
}) as const,
)
let lastUsage: CompletionUsage | undefined
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
for (const matcherChunk of matcher.update(delta.content)) {
yield matcherChunk
}
}
if (chunk.usage) {
lastUsage = chunk.usage
}
}
for (const chunk of matcher.final()) {
yield chunk
}
if (lastUsage) {
yield {
type: "usage",
inputTokens: lastUsage?.prompt_tokens || 0,
outputTokens: lastUsage?.completion_tokens || 0,
}
}
}
override getModel(): { id: string; info: ModelInfo } {
return {
id: this.options.ollamaModelId || "",
info: openAiModelInfoSaneDefaults,
}
}
async completePrompt(prompt: string): Promise<string> {
try {
const modelId = this.getModel().id
const useR1Format = modelId.toLowerCase().includes("deepseek-r1")
let response
try {
response = await this.client.chat.completions.create({
model: this.getModel().id,
messages: useR1Format
? convertToR1Format([{ role: "user", content: prompt }])
: [{ role: "user", content: prompt }],
temperature: this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
stream: false,
})
} catch (error) {
throw handleOpenAIError(error, this.providerName)
}
return response.choices[0]?.message.content || ""
} catch (error) {
if (error instanceof Error) {
throw new Error(`Ollama completion error: ${error.message}`)
}
throw error
}
}
}