From ce1ddc93255b527da4bd097af8bb238417a96725 Mon Sep 17 00:00:00 2001 From: Pablo Fernandez Date: Mon, 25 May 2026 14:17:56 +0000 Subject: [PATCH 1/3] feat(ollama): support tool_choice in completion requests The Ollama /api/chat endpoint accepts tool_choice (none/auto/required) in versions that support tool calling. Previously rig warned and dropped this field, which meant the rig Extractor's tool_choice: Required was silently ignored, forcing reliance on prompt instructions alone to produce structured output. This commit: - Adds OllamaToolChoice enum mapping rig's ToolChoice variants - Wires tool_choice through OllamaCompletionRequest serialization - Removes the tracing::warn that fired on every extractor call - Adds unit tests verifying correct serialization for all variants --- crates/rig-core/src/providers/ollama.rs | 138 +++++++++++++++++++++++- 1 file changed, 135 insertions(+), 3 deletions(-) diff --git a/crates/rig-core/src/providers/ollama.rs b/crates/rig-core/src/providers/ollama.rs index 614bff03b..3d5514442 100644 --- a/crates/rig-core/src/providers/ollama.rs +++ b/crates/rig-core/src/providers/ollama.rs @@ -418,6 +418,29 @@ impl TryFrom for completion::CompletionResponse for ToolChoice { + type Error = CompletionError; + fn try_from(value: message::ToolChoice) -> Result { + match value { + message::ToolChoice::Specific { .. } => Err(CompletionError::ProviderError( + "Ollama does not support specific tool choice".to_string(), + )), + message::ToolChoice::Auto => Ok(Self::Auto), + message::ToolChoice::None => Ok(Self::None), + message::ToolChoice::Required => Ok(Self::Required), + } + } +} + #[derive(Debug, Serialize, Deserialize)] pub(super) struct OllamaCompletionRequest { model: String, @@ -434,6 +457,8 @@ pub(super) struct OllamaCompletionRequest { keep_alive: Option, #[serde(skip_serializing_if = "Option::is_none")] format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tool_choice: Option, options: serde_json::Value, } @@ -442,9 +467,6 @@ impl TryFrom<(&str, CompletionRequest)> for OllamaCompletionRequest { fn try_from((model, req): (&str, CompletionRequest)) -> Result { let model = req.model.clone().unwrap_or_else(|| model.to_string()); - if req.tool_choice.is_some() { - tracing::warn!("WARNING: `tool_choice` not supported for Ollama"); - } // Build up the order of messages (context, chat_history, prompt) let mut partial_history = vec![]; if let Some(docs) = req.normalized_documents() { @@ -502,6 +524,8 @@ impl TryFrom<(&str, CompletionRequest)> for OllamaCompletionRequest { json!({ "temperature": req.temperature }) }; + let tool_choice = req.tool_choice.map(ToolChoice::try_from).transpose()?; + Ok(Self { model: model.to_string(), messages: full_history, @@ -511,6 +535,7 @@ impl TryFrom<(&str, CompletionRequest)> for OllamaCompletionRequest { think, keep_alive, format: req.output_schema, + tool_choice, tools: req .tools .clone() @@ -1766,6 +1791,113 @@ mod tests { ); } + #[test] + fn test_completion_request_with_tool_choice_required() { + use crate::OneOrMany; + use crate::completion::Message as CompletionMessage; + use crate::message::{Text, ToolChoice as MessageToolChoice, UserContent}; + + let completion_request = CompletionRequest { + model: Some("llama3.1".to_string()), + preamble: None, + chat_history: OneOrMany::one(CompletionMessage::User { + content: OneOrMany::one(UserContent::Text(Text { + text: "Hello!".to_string(), + })), + }), + documents: vec![], + tools: vec![], + temperature: None, + max_tokens: None, + tool_choice: Some(MessageToolChoice::Required), + additional_params: None, + output_schema: None, + }; + + let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request)) + .expect("Failed to create Ollama request"); + + let serialized = + serde_json::to_value(&ollama_request).expect("Failed to serialize request"); + + assert_eq!( + serialized.get("tool_choice").and_then(|v| v.as_str()), + Some("required"), + "tool_choice should serialize to 'required'" + ); + } + + #[test] + fn test_completion_request_without_tool_choice() { + use crate::OneOrMany; + use crate::completion::Message as CompletionMessage; + use crate::message::{Text, UserContent}; + + let completion_request = CompletionRequest { + model: Some("llama3.1".to_string()), + preamble: None, + chat_history: OneOrMany::one(CompletionMessage::User { + content: OneOrMany::one(UserContent::Text(Text { + text: "Hello!".to_string(), + })), + }), + documents: vec![], + tools: vec![], + temperature: None, + max_tokens: None, + tool_choice: None, + additional_params: None, + output_schema: None, + }; + + let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request)) + .expect("Failed to create Ollama request"); + + let serialized = + serde_json::to_value(&ollama_request).expect("Failed to serialize request"); + + assert!( + serialized.get("tool_choice").is_none(), + "tool_choice field should be absent when None" + ); + } + + #[test] + fn test_completion_request_with_tool_choice_auto() { + use crate::OneOrMany; + use crate::completion::Message as CompletionMessage; + use crate::message::{Text, ToolChoice as MessageToolChoice, UserContent}; + + let completion_request = CompletionRequest { + model: Some("llama3.1".to_string()), + preamble: None, + chat_history: OneOrMany::one(CompletionMessage::User { + content: OneOrMany::one(UserContent::Text(Text { + text: "Hello!".to_string(), + })), + }), + documents: vec![], + tools: vec![], + temperature: None, + max_tokens: None, + tool_choice: Some(MessageToolChoice::Auto), + additional_params: None, + output_schema: None, + }; + + let ollama_request = OllamaCompletionRequest::try_from(("llama3.1", completion_request)) + .expect("Failed to create Ollama request"); + + let serialized = + serde_json::to_value(&ollama_request).expect("Failed to serialize request"); + + assert_eq!( + serialized.get("tool_choice").and_then(|v| v.as_str()), + Some("auto"), + "tool_choice should serialize to 'auto'" + ); + } + #[test] fn test_client_initialization() { let _client = crate::providers::ollama::Client::new(Nothing).expect("Client::new() failed"); From bdf9c1ac2a3f113e6e4e92701dd067c1c5958112 Mon Sep 17 00:00:00 2001 From: Pablo Fernandez Date: Mon, 25 May 2026 15:14:13 +0000 Subject: [PATCH 2/3] restore lib name to 'rig' for compatibility with 0.35 consumers --- crates/rig-core/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/rig-core/Cargo.toml b/crates/rig-core/Cargo.toml index 81caf13c5..977534c6e 100644 --- a/crates/rig-core/Cargo.toml +++ b/crates/rig-core/Cargo.toml @@ -15,7 +15,7 @@ rustdoc-args = ["--cfg", "docsrs"] workspace = true [lib] -name = "rig_core" +name = "rig" path = "src/lib.rs" doctest = true From b5179b0e3ed52a92d40f2c091cc9d8a6d34795f9 Mon Sep 17 00:00:00 2001 From: Pablo Fernandez Date: Mon, 25 May 2026 19:35:20 +0000 Subject: [PATCH 3/3] revert: restore rig_core lib name (correct upstream name) --- crates/rig-core/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/rig-core/Cargo.toml b/crates/rig-core/Cargo.toml index 977534c6e..81caf13c5 100644 --- a/crates/rig-core/Cargo.toml +++ b/crates/rig-core/Cargo.toml @@ -15,7 +15,7 @@ rustdoc-args = ["--cfg", "docsrs"] workspace = true [lib] -name = "rig" +name = "rig_core" path = "src/lib.rs" doctest = true