Skip to content

feat: ollama tool shim #1448

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alicehau merged 30 commits into from
Mar 12, 2025
Merged

feat: ollama tool shim #1448

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
e1b5ae9
ollama model + structured output shim
alicehau Feb 27, 2025
220f926
tool shim generic
alicehau Feb 28, 2025
b9b6fe8
refactoring toolshim
alicehau Feb 28, 2025
d666bcb
refactor
alicehau Feb 28, 2025
73407cd
ollama toolshim ready
alicehau Feb 28, 2025
05530b3
clean up
alicehau Feb 28, 2025
bba2741
lint
alicehau Mar 4, 2025
acfbf83
format
alicehau Mar 4, 2025
48544e3
add toolshim to openai provider
alicehau Mar 4, 2025
92122b6
update toolshim prompt:
Mar 5, 2025
79aa25c
prompt update
alicehau Mar 5, 2025
50a1a5f
refactor
Mar 7, 2025
dc67538
undo change to reference agent
Mar 7, 2025
4b8805d
undo ref agent changes
Mar 10, 2025
004d032
refactor config
Mar 10, 2025
6877687
refactor
Mar 10, 2025
5b035cb
check env vars
Mar 10, 2025
23a00e3
fix tests
Mar 10, 2025
bbdbe66
clean up
Mar 10, 2025
9ad5d3a
start moving interpret chat tool check into agent
Mar 10, 2025
8f76da0
fully move shim logic into agent
Mar 10, 2025
374ede5
minor semicolon fix
Mar 10, 2025
255203d
refactor toolshim logic out of providers
Mar 10, 2025
190a26c
revert providers and rebase
Mar 10, 2025
ed61236
refactor
Mar 10, 2025
e6866c5
lint
Mar 10, 2025
b22340d
Merge branch 'main' into ahau/toolshim-refactor
alicehau Mar 12, 2025
f7f71ee
rename var
Mar 12, 2025
6de0be6
rename vars to toolshim
Mar 12, 2025
f0bc9f6
renaming
Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions crates/goose-cli/src/commands/configure.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,8 +296,16 @@ pub async fn configure_provider_dialog() -> Result<bool, Box<dyn Error>> {
let spin = spinner();
spin.start("Checking your configuration...");

// Use max tokens to speed up the provider test.
let model_config = goose::model::ModelConfig::new(model.clone()).with_max_tokens(Some(50));
// Create model config with env var settings
let model_config = goose::model::ModelConfig::new(model.clone())
.with_max_tokens(Some(50))
.with_toolshim(
std::env::var("GOOSE_TOOLSHIM")
.map(|val| val == "1" || val.to_lowercase() == "true")
.unwrap_or(false),
)
.with_toolshim_model(std::env::var("GOOSE_TOOLSHIM_OLLAMA_MODEL").ok());

let provider = create(provider_name, model_config)?;

let messages =
Expand Down
25 changes: 23 additions & 2 deletions crates/goose/src/agents/truncate.rs
View file
Open in desktop
Copy link
Contributor

@salman1993 salman1993 Mar 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we also need to add to reference and summarize agents?

Copy link
Contributor

@ahau-square ahau-square Mar 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thought those were just legacy at this point? I don't recall any way for folks to switch their agents?

Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ use crate::message::{Message, ToolRequest};
use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
use crate::providers::errors::ProviderError;
use crate::providers::toolshim::{
augment_message_with_tool_calls, modify_system_prompt_for_tool_json, OllamaInterpreter,
};
use crate::register_agent;
use crate::session;
use crate::token_counter::TokenCounter;
Expand Down Expand Up @@ -217,7 +220,17 @@ impl Agent for TruncateAgent {
tools.push(list_resources_tool);
}

let system_prompt = capabilities.get_system_prompt().await;
let config = capabilities.provider().get_model_config();
let mut system_prompt = capabilities.get_system_prompt().await;
let mut toolshim_tools = vec![];
if config.toolshim {
// If tool interpretation is enabled, modify the system prompt to instruct to return JSON tool requests
system_prompt = modify_system_prompt_for_tool_json(&system_prompt, &tools);
// make a copy of tools before empty
toolshim_tools = tools.clone();
// pass empty tools vector to provider completion since toolshim will handle tool calls instead
tools = vec![];
}

// Set the user_message field in the span instead of creating a new event
if let Some(content) = messages
Expand All @@ -236,7 +249,15 @@ impl Agent for TruncateAgent {
&messages,
&tools,
).await {
Ok((response, usage)) => {
Ok((mut response, usage)) => {
// Post-process / structure the response only if tool interpretation is enabled
if config.toolshim {
let interpreter = OllamaInterpreter::new()
.map_err(|e| anyhow::anyhow!("Failed to create OllamaInterpreter: {}", e))?;

response = augment_message_with_tool_calls(&interpreter, response, &toolshim_tools).await?;
}

capabilities.record_usage(usage.clone()).await;

// record usage for the session in the session file
Expand Down
42 changes: 41 additions & 1 deletion crates/goose/src/model.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ pub struct ModelConfig {
pub temperature: Option<f32>,
/// Optional maximum tokens to generate
pub max_tokens: Option<i32>,
/// Whether to interpret tool calls with toolshim
pub toolshim: bool,
/// Model to use for toolshim (optional as a default exists)
pub toolshim_model: Option<String>,
}

impl ModelConfig {
Expand All @@ -34,12 +38,20 @@ impl ModelConfig {
let context_limit = Self::get_model_specific_limit(&model_name);
let tokenizer_name = Self::infer_tokenizer_name(&model_name);

let toolshim = std::env::var("GOOSE_TOOLSHIM")
.map(|val| val == "1" || val.to_lowercase() == "true")
.unwrap_or(false);

let toolshim_model = std::env::var("GOOSE_TOOLSHIM_OLLAMA_MODEL").ok();

Self {
model_name,
tokenizer_name: tokenizer_name.to_string(),
context_limit,
temperature: None,
max_tokens: None,
toolshim,
toolshim_model,
}
}

Expand Down Expand Up @@ -96,7 +108,19 @@ impl ModelConfig {
self
}

// Get the tokenizer name
/// Set whether to interpret tool calls
pub fn with_toolshim(mut self, toolshim: bool) -> Self {
self.toolshim = toolshim;
self
}

/// Set the tool call interpreter model
pub fn with_toolshim_model(mut self, model: Option<String>) -> Self {
self.toolshim_model = model;
self
}

/// Get the tokenizer name
pub fn tokenizer_name(&self) -> &str {
&self.tokenizer_name
}
Expand Down Expand Up @@ -142,4 +166,20 @@ mod tests {
assert_eq!(config.max_tokens, Some(1000));
assert_eq!(config.context_limit, Some(50_000));
}

#[test]
fn test_model_config_tool_interpretation() {
// Test without env vars - should be false
let config = ModelConfig::new("test-model".to_string());
assert!(!config.toolshim);

// Test with tool interpretation setting
let config = ModelConfig::new("test-model".to_string()).with_toolshim(true);
assert!(config.toolshim);

// Test tool interpreter model
let config = ModelConfig::new("test-model".to_string())
.with_toolshim_model(Some("mistral-nemo".to_string()));
assert_eq!(config.toolshim_model, Some("mistral-nemo".to_string()));
}
}
6 changes: 6 additions & 0 deletions crates/goose/src/providers/formats/openai.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -826,6 +826,8 @@ mod tests {
context_limit: Some(4096),
temperature: None,
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
let obj = request.as_object().unwrap();
Expand Down Expand Up @@ -856,6 +858,8 @@ mod tests {
context_limit: Some(4096),
temperature: None,
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
let obj = request.as_object().unwrap();
Expand Down Expand Up @@ -887,6 +891,8 @@ mod tests {
context_limit: Some(4096),
temperature: None,
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
let obj = request.as_object().unwrap();
Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/providers/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ pub mod oauth;
pub mod ollama;
pub mod openai;
pub mod openrouter;
pub mod toolshim;
pub mod utils;

pub use factory::{create, providers};
87 changes: 12 additions & 75 deletions crates/goose/src/providers/ollama.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use crate::model::ModelConfig;
use crate::providers::formats::openai::{create_request, get_usage, response_to_message};
use anyhow::Result;
use async_trait::async_trait;
use indoc::formatdoc;
use mcp_core::tool::Tool;
use reqwest::Client;
use serde_json::Value;
Expand Down Expand Up @@ -53,8 +52,8 @@ impl OllamaProvider {
})
}

async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
// TODO: remove this later when the UI handles provider config refresh
/// Get the base URL for Ollama API calls
fn get_base_url(&self) -> Result<Url, ProviderError> {
// OLLAMA_HOST is sometimes just the 'host' or 'host:port' without a scheme
let base = if self.host.starts_with("http://") || self.host.starts_with("https://") {
self.host.clone()
Expand All @@ -73,6 +72,13 @@ impl OllamaProvider {
})?;
}

Ok(base_url)
}

async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
// TODO: remove this later when the UI handles provider config refresh
let base_url = self.get_base_url()?;

let url = base_url.join("v1/chat/completions").map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
Expand Down Expand Up @@ -116,86 +122,17 @@ impl Provider for OllamaProvider {
messages: &[Message],
tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
// Transform the system message to replace developer instructions
let modified_system = if let Some(dev_section) = system.split("## developer").nth(1) {
if let (Some(start_idx), Some(end_idx)) = (
dev_section.find("### Instructions"),
dev_section.find("operating system:"),
) {
let new_instructions = formatdoc! {r#"
The Developer extension enables you to edit code files, execute shell commands, and capture screen/window content. These tools allow for various development and debugging workflows.
Available Tools:
1. Shell Execution (`shell`)
Executes commands in the shell and returns the combined output and error messages.
Use cases:
- Running scripts: `python script.py`
- Installing dependencies: `pip install -r requirements.txt`
- Checking system information: `uname -a`, `df -h`
- Searching for files or text: **Use `rg` (ripgrep) instead of `find` or `ls -r`**
- Find a file: `rg --files | rg example.py`
- Search within files: `rg 'class Example'`
Best Practices:
- **Avoid commands with large output** (pipe them to a file if necessary).
- **Run background processes** if they take a long time (e.g., `uvicorn main:app &`).
- **git commands can be run on the shell, however if the git extension is installed, you should use the git tool instead.
- **If the shell command is a rm, mv, or cp, you should verify with the user before running the command.
2. Text Editor (`text_editor`)
Performs file-based operations such as viewing, writing, replacing text, and undoing edits.
Commands:
- view: Read the content of a file.
- write: Create or overwrite a file. Caution: Overwrites the entire file!
- str_replace: Replace a specific string in a file.
- undo_edit: Revert the last edit.
Example Usage:
text_editor(command="view", file_path="/absolute/path/to/file.py")
text_editor(command="write", file_path="/absolute/path/to/file.py", file_text="print('hello world')")
text_editor(command="str_replace", file_path="/absolute/path/to/file.py", old_str="hello world", new_str="goodbye world")
text_editor(command="undo_edit", file_path="/absolute/path/to/file.py")
Protocol for Text Editor:
For edit and replace commands, please verify what you are editing with the user before running the command.
- User: "Please edit the file /absolute/path/to/file.py"
- Assistant: "Ok sounds good, I'll be editing the file /absolute/path/to/file.py and creating modifications xyz to the file. Let me know whether you'd like to proceed."
- User: "Yes, please proceed."
- Assistant: "I've created the modifications xyz to the file /absolute/path/to/file.py"
3. List Windows (`list_windows`)
Lists all visible windows with their titles.
Use this to find window titles for screen capture.
4. Screen Capture (`screen_capture`)
Takes a screenshot of a display or specific window.
Options:
- Capture display: `screen_capture(display=0)` # Main display
- Capture window: `screen_capture(window_title="Window Title")`
Info: at the start of the session, the user's directory is:
"#};

let before_dev = system.split("## developer").next().unwrap_or("");
let after_marker = &dev_section[end_idx..];

format!(
"{}## developer{}### Instructions\n{}{}",
before_dev,
&dev_section[..start_idx],
new_instructions,
after_marker
)
} else {
system.to_string()
}
} else {
system.to_string()
};

let payload = create_request(
&self.model,
&modified_system,
system,
messages,
tools,
&super::utils::ImageFormat::OpenAi,
)?;
let response = self.post(payload.clone()).await?;

// Parse response
let response = self.post(payload.clone()).await?;
let message = response_to_message(response.clone())?;

let usage = match get_usage(&response) {
Ok(usage) => usage,
Err(ProviderError::UsageError(e)) => {
Expand Down
Loading
Loading