Skip to content

feat: Implement streaming for OpenAI #3413

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
jamadeo merged 2 commits into from
Jul 15, 2025
Merged

feat: Implement streaming for OpenAI #3413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bb58460
Implement streaming for OpenAI
jamadeo Jul 2, 2025
2ca4466
Clean up duplicated code and get_usage
jamadeo Jul 14, 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: 4 additions & 8 deletions crates/goose/src/providers/azure.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -249,14 +249,10 @@ impl Provider for AzureProvider {
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)) => {
tracing::debug!("Failed to get usage data: {}", e);
Usage::default()
}
Err(e) => return Err(e),
};
let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
tracing::debug!("Failed to get usage data");
Usage::default()
});
let model = get_model(&response);
emit_debug_trace(&self.model, &payload, &response, &usage);
Ok((message, ProviderUsage::new(model, usage)))
Expand Down
15 changes: 6 additions & 9 deletions crates/goose/src/providers/databricks.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,13 @@ use tokio_util::io::StreamReader;
use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::embedding::EmbeddingCapable;
use super::errors::ProviderError;
use super::formats::databricks::{create_request, get_usage, response_to_message};
use super::formats::databricks::{create_request, response_to_message};
use super::oauth;
use super::utils::{get_model, ImageFormat};
use crate::config::ConfigError;
use crate::message::Message;
use crate::model::ModelConfig;
use crate::providers::formats::databricks::response_to_streaming_message;
use crate::providers::formats::openai::{get_usage, response_to_streaming_message};
use mcp_core::tool::Tool;
use serde_json::json;
use tokio::time::sleep;
Expand Down Expand Up @@ -455,13 +455,10 @@ impl Provider for DatabricksProvider {

// Parse response
let message = response_to_message(response.clone())?;
let usage = match response.get("usage").map(get_usage) {
Some(usage) => usage,
None => {
tracing::debug!("Failed to get usage data");
Usage::default()
}
};
let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
tracing::debug!("Failed to get usage data");
Usage::default()
});
let model = get_model(&response);
super::utils::emit_debug_trace(&self.model, &payload, &response, &usage);

Expand Down
137 changes: 0 additions & 137 deletions crates/goose/src/providers/formats/databricks.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,10 @@
use crate::message::{Message, MessageContent};
use crate::model::ModelConfig;
use crate::providers::base::{ProviderUsage, Usage};
use crate::providers::utils::{
convert_image, detect_image_path, is_valid_function_name, load_image_file,
sanitize_function_name, ImageFormat,
};
use anyhow::{anyhow, Error};
use async_stream::try_stream;
use futures::Stream;
use mcp_core::ToolError;
use mcp_core::{Content, Role, Tool, ToolCall};
use serde::{Deserialize, Serialize};
Expand Down Expand Up @@ -404,140 +401,6 @@ struct StreamingChunk {
model: String,
}

fn strip_data_prefix(line: &str) -> Option<&str> {
line.strip_prefix("data: ").map(|s| s.trim())
}

pub fn response_to_streaming_message<S>(
mut stream: S,
) -> impl Stream<Item = anyhow::Result<(Option<Message>, Option<ProviderUsage>)>> + 'static
where
S: Stream<Item = anyhow::Result<String>> + Unpin + Send + 'static,
{
try_stream! {
use futures::StreamExt;

'outer: while let Some(response) = stream.next().await {
if response.as_ref().is_ok_and(|s| s == "data: [DONE]") {
break 'outer;
}
let response_str = response?;
let line = strip_data_prefix(&response_str);

if line.is_none() || line.is_some_and(|l| l.is_empty()) {
continue
}

let chunk: StreamingChunk = serde_json::from_str(line
.ok_or_else(|| anyhow!("unexpected stream format"))?)
.map_err(|e| anyhow!("Failed to parse streaming chunk: {}: {:?}", e, &line))?;
let model = chunk.model.clone();

let usage = chunk.usage.as_ref().map(|u| {
ProviderUsage {
usage: get_usage(u),
model,
}
});

if chunk.choices.is_empty() {
yield (None, usage)
} else if let Some(tool_calls) = &chunk.choices[0].delta.tool_calls {
let tool_call = &tool_calls[0];
let id = tool_call.id.clone().ok_or(anyhow!("No tool call ID"))?;
let function_name = tool_call.function.name.clone().ok_or(anyhow!("No function name"))?;
let mut arguments = tool_call.function.arguments.clone();

while let Some(response_chunk) = stream.next().await {
if response_chunk.as_ref().is_ok_and(|s| s == "data: [DONE]") {
break 'outer;
}
let response_str = response_chunk?;
if let Some(line) = strip_data_prefix(&response_str) {
let tool_chunk: StreamingChunk = serde_json::from_str(line)
.map_err(|e| anyhow!("Failed to parse streaming chunk: {}: {:?}", e, &line))?;
let more_args = tool_chunk.choices[0].delta.tool_calls.as_ref()
.and_then(|calls| calls.first())
.map(|call| call.function.arguments.as_str());
if let Some(more_args) = more_args {
arguments.push_str(more_args);
} else {
break;
}
}
}

let parsed = if arguments.is_empty() {
Ok(json!({}))
} else {
serde_json::from_str::<Value>(&arguments)
};

let content = match parsed {
Ok(params) => MessageContent::tool_request(
id,
Ok(ToolCall::new(function_name, params)),
),
Err(e) => {
let error = ToolError::InvalidParameters(format!(
"Could not interpret tool use parameters for id {}: {}",
id, e
));
MessageContent::tool_request(id, Err(error))
}
};

yield (
Some(Message {
id: chunk.id,
role: Role::Assistant,
created: chrono::Utc::now().timestamp(),
content: vec![content],
}),
usage,
)
} else if let Some(text) = &chunk.choices[0].delta.content {
yield (
Some(Message {
id: chunk.id,
role: Role::Assistant,
created: chrono::Utc::now().timestamp(),
content: vec![MessageContent::text(text)],
}),
if chunk.choices[0].finish_reason.is_some() {
usage
} else {
None
},
)
}
}
}
}

pub fn get_usage(usage: &Value) -> Usage {
let input_tokens = usage
.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|v| v as i32);

let output_tokens = usage
.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|v| v as i32);

let total_tokens = usage
.get("total_tokens")
.and_then(|v| v.as_i64())
.map(|v| v as i32)
.or_else(|| match (input_tokens, output_tokens) {
(Some(input), Some(output)) => Some(input + output),
_ => None,
});

Usage::new(input_tokens, output_tokens, total_tokens)
}

/// Validates and fixes tool schemas to ensure they have proper parameter structure.
/// If parameters exist, ensures they have properties and required fields, or removes parameters entirely.
pub fn validate_tool_schemas(tools: &mut [Value]) {
Expand Down
Loading
Loading