From 10dfe06c6529dae38602dbeaa4dd6c633597efc6 Mon Sep 17 00:00:00 2001 From: jovonni Date: Fri, 22 Aug 2025 21:44:44 -0400 Subject: [PATCH 1/2] stream for ollama Signed-off-by: jovonni --- crates/goose/src/providers/ollama.rs | 58 ++++++++++++++++++++++++++-- 1 file changed, 54 insertions(+), 4 deletions(-) diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index 1479a847fe73..7951052c5dc3 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -1,21 +1,28 @@ use super::api_client::{ApiClient, AuthMethod}; -use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; +use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata, ProviderUsage, Usage}; use super::errors::ProviderError; use super::retry::ProviderRetry; -use super::utils::{get_model, handle_response_openai_compat}; +use super::utils::{get_model, handle_response_openai_compat, handle_status_openai_compat}; use crate::config::custom_providers::CustomProviderConfig; use crate::conversation::message::Message; use crate::conversation::Conversation; use crate::impl_provider_default; use crate::model::ModelConfig; -use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; +use crate::providers::formats::openai::{create_request, get_usage, response_to_message, response_to_streaming_message}; use crate::utils::safe_truncate; use anyhow::Result; +use async_stream::try_stream; use async_trait::async_trait; +use futures::TryStreamExt; use regex::Regex; use rmcp::model::Tool; -use serde_json::Value; +use serde_json::{json, Value}; +use std::io; use std::time::Duration; +use tokio::pin; +use tokio_stream::StreamExt; +use tokio_util::codec::{FramedRead, LinesCodec}; +use tokio_util::io::StreamReader; use url::Url; pub const OLLAMA_HOST: &str = "localhost"; @@ -228,6 +235,49 @@ impl Provider for OllamaProvider { fn supports_streaming(&self) -> bool { self.supports_streaming } + + async fn stream( + &self, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let config = crate::config::Config::global(); + let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string()); + let filtered_tools = if goose_mode == "chat" { &[] } else { tools }; + + let mut payload = create_request( + &self.model, + system, + messages, + filtered_tools, + &super::utils::ImageFormat::OpenAi, + )?; + payload["stream"] = json!(true); + payload["stream_options"] = json!({ + "include_usage": true, + }); + + let response = self + .api_client + .response_post("v1/chat/completions", &payload) + .await?; + let response = handle_status_openai_compat(response).await?; + let stream = response.bytes_stream().map_err(io::Error::other); + let model_config = self.model.clone(); + + Ok(Box::pin(try_stream! { + let stream_reader = StreamReader::new(stream); + let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from); + let message_stream = response_to_streaming_message(framed); + pin!(message_stream); + while let Some(message) = message_stream.next().await { + let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?; + super::utils::emit_debug_trace(&model_config, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default()); + yield (message, usage); + } + })) + } } impl OllamaProvider { From 31ec3414792ba954791aef290c8ec684a77dd191 Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Wed, 27 Aug 2025 11:00:57 +1000 Subject: [PATCH 2/2] streaming works, removed tool check unneeded Signed-off-by: jovonni --- crates/goose/src/providers/ollama.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index 7951052c5dc3..95c57a65c414 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -8,7 +8,9 @@ use crate::conversation::message::Message; use crate::conversation::Conversation; use crate::impl_provider_default; use crate::model::ModelConfig; -use crate::providers::formats::openai::{create_request, get_usage, response_to_message, response_to_streaming_message}; +use crate::providers::formats::openai::{ + create_request, get_usage, response_to_message, response_to_streaming_message, +}; use crate::utils::safe_truncate; use anyhow::Result; use async_stream::try_stream; @@ -85,7 +87,7 @@ impl OllamaProvider { Ok(Self { api_client, model, - supports_streaming: false, + supports_streaming: true, }) } @@ -242,15 +244,11 @@ impl Provider for OllamaProvider { messages: &[Message], tools: &[Tool], ) -> Result { - let config = crate::config::Config::global(); - let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string()); - let filtered_tools = if goose_mode == "chat" { &[] } else { tools }; - let mut payload = create_request( &self.model, system, messages, - filtered_tools, + tools, &super::utils::ImageFormat::OpenAi, )?; payload["stream"] = json!(true);