Produce Thought/ThoughtChunk for Fireworks reasoning_content (tensorzero#3006)

* Produce Thought/ThoughtChunk for Fireworks `reasoning_content`

* Fix clippy

Author: Aaron1011

tensorzero-core/src/providers/fireworks/mod.rs

@@ -1,6 +1,6 @@
 use std::{borrow::Cow, sync::OnceLock};
 
-use crate::tool::Tool;
+use crate::{providers::helpers_thinking_block::THINK_CHUNK_ID, tool::Tool};
 use futures::StreamExt;
 use lazy_static::lazy_static;
 use reqwest_eventsource::{Event, EventSource};
@@ -470,6 +470,8 @@ struct FireworksResponseMessage {
     #[serde(skip_serializing_if = "Option::is_none")]
     content: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    reasoning_content: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     tool_calls: Option<Vec<FireworksResponseToolCall>>,
 }
 
@@ -531,6 +533,8 @@ struct FireworksDelta {
     #[serde(skip_serializing_if = "Option::is_none")]
     content: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    reasoning_content: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     tool_calls: Option<Vec<FireworksToolCallChunk>>,
 }
 
@@ -630,6 +634,14 @@ fn fireworks_to_tensorzero_chunk(
         if let Some(reason) = choice.finish_reason {
             finish_reason = Some(reason.into());
         }
+        if let Some(reasoning) = choice.delta.reasoning_content {
+            content.push(ContentBlockChunk::Thought(ThoughtChunk {
+                text: Some(reasoning),
+                signature: None,
+                id: THINK_CHUNK_ID.to_string(),
+                provider_type: Some(PROVIDER_TYPE.to_string()),
+            }));
+        }
         if let Some(text) = choice.delta.content {
             if parse_think_blocks {
                 if !thinking_state.update(&text, PROVIDER_TYPE)? {
@@ -745,6 +757,13 @@ impl<'a> TryFrom<FireworksResponseWithMetadata<'a>> for ProviderInferenceRespons
             }
             ))?;
         let mut content: Vec<ContentBlockOutput> = Vec::new();
+        if let Some(reasoning) = message.reasoning_content {
+            content.push(ContentBlockOutput::Thought(Thought {
+                text: Some(reasoning),
+                signature: None,
+                provider_type: Some(PROVIDER_TYPE.to_string()),
+            }));
+        }
         if let Some(raw_text) = message.content {
             let (clean_text, extracted_reasoning) =
                 process_think_blocks(&raw_text, parse_think_blocks, PROVIDER_TYPE)?;
@@ -826,6 +845,7 @@ mod tests {
                 index: 0,
                 finish_reason: Some(FireworksFinishReason::Stop),
                 message: FireworksResponseMessage {
+                    reasoning_content: None,
                     content: Some(test_response_with_thinking.to_string()),
                     tool_calls: None,
                 },
@@ -1008,6 +1028,7 @@ mod tests {
                 index: 0,
                 finish_reason: Some(FireworksFinishReason::Stop),
                 message: FireworksResponseMessage {
+                    reasoning_content: None,
                     content: Some("Hello, world!".to_string()),
                     tool_calls: None,
                 },
@@ -1077,6 +1098,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: Some("Hello".to_string()),
+                    reasoning_content: None,
                     tool_calls: None,
                 },
                 finish_reason: Some(FireworksFinishReason::Stop),
@@ -1109,6 +1131,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: None,
+                    reasoning_content: None,
                     tool_calls: Some(vec![FireworksToolCallChunk {
                         index: 0,
                         id: None,
@@ -1178,6 +1201,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: Some("<think>".to_string()),
+                    reasoning_content: None,
                     tool_calls: None,
                 },
                 finish_reason: None,
@@ -1209,6 +1233,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: Some("reasoning".to_string()),
+                    reasoning_content: None,
                     tool_calls: None,
                 },
                 finish_reason: None,
@@ -1241,6 +1266,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: Some("</think>".to_string()),
+                    reasoning_content: None,
                     tool_calls: None,
                 },
                 finish_reason: None,
@@ -1268,6 +1294,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: Some("Final answer".to_string()),
+                    reasoning_content: None,
                     tool_calls: None,
                 },
                 finish_reason: None,
@@ -1302,6 +1329,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: Some("Hello <think>should not parse</think>".to_string()),
+                    reasoning_content: None,
                     tool_calls: None,
                 },
                 finish_reason: Some(FireworksFinishReason::Stop),
@@ -1335,6 +1363,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: None,
+                    reasoning_content: None,
                     tool_calls: Some(vec![FireworksToolCallChunk {
                         index: 0,
                         id: Some("new_id".to_string()),
@@ -1379,6 +1408,7 @@ mod tests {
             choices: vec![FireworksChatChunkChoice {
                 delta: FireworksDelta {
                     content: None,
+                    reasoning_content: None,
                     tool_calls: Some(vec![FireworksToolCallChunk {
                         index: 0,
                         id: None,

tensorzero-core/src/providers/helpers_thinking_block.rs

@@ -58,6 +58,8 @@ pub enum ThinkingState {
     Finished,
 }
 
+pub const THINK_CHUNK_ID: u64 = 1;
+
 impl ThinkingState {
     pub fn get_id(&self) -> String {
         match self {

tensorzero-core/tests/e2e/providers/fireworks.rs

@@ -1,6 +1,15 @@
+#![expect(clippy::print_stderr)]
 use std::collections::HashMap;
 
-use crate::providers::common::{E2ETestProvider, E2ETestProviders};
+use http::StatusCode;
+use reqwest::Client;
+use reqwest_eventsource::{Event, RequestBuilderExt};
+use tokio_stream::StreamExt;
+
+use crate::{
+    common::get_gateway_endpoint,
+    providers::common::{E2ETestProvider, E2ETestProviders},
+};
 
 crate::generate_provider_tests!(get_providers);
 crate::generate_batch_inference_tests!(get_providers);
@@ -117,3 +126,81 @@ async fn get_providers() -> E2ETestProviders {
         shorthand_inference: shorthand_providers,
     }
 }
+
+#[tokio::test]
+async fn test_fireworks_reasoning_content_non_stream() {
+    let response = Client::new()
+        .post(get_gateway_endpoint("/inference"))
+        .json(&serde_json::json!({
+            "model_name": "gpt-oss-20b-fireworks",
+            "input": {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "What is the capital of France? Think before responding."
+                    }
+                ]
+            }
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let response_body = response.json::<serde_json::Value>().await.unwrap();
+
+    eprintln!("API response: {response_body}");
+    let content = response_body["content"].as_array().unwrap();
+    // Check that the response contains a thought block
+    let thought = content.iter().find(|c| c["type"] == "thought").unwrap();
+    assert!(
+        !thought["text"].as_str().unwrap().is_empty(),
+        "Thought block was empty: {thought:?}",
+    );
+}
+
+#[tokio::test]
+async fn test_fireworks_reasoning_content_stream() {
+    let mut event_source = Client::new()
+        .post(get_gateway_endpoint("/inference"))
+        .json(&serde_json::json!({
+            "model_name": "gpt-oss-20b-fireworks",
+            "input": {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": "What is the capital of France? Think before responding."
+                    }
+                ]
+            },
+            "stream": true,
+        }))
+        .eventsource()
+        .unwrap();
+
+    let mut chunks = vec![];
+    while let Some(event) = event_source.next().await {
+        let event = event.unwrap();
+        match event {
+            Event::Open => continue,
+            Event::Message(message) => {
+                eprintln!("API chunk: {message:?}");
+                if message.data == "[DONE]" {
+                    break;
+                }
+                chunks.push(message.data.to_string());
+            }
+        }
+    }
+
+    let mut found_thought = false;
+    for chunk in chunks {
+        let chunk_json: serde_json::Value = serde_json::from_str(&chunk).unwrap();
+        eprintln!("Chunk: {chunk_json}");
+        if chunk_json["content"][0]["type"] == "thought" {
+            found_thought = true;
+            break;
+        }
+    }
+    assert!(found_thought, "Thought chunk not found");
+}

tensorzero-core/tests/e2e/tensorzero.toml

@@ -39,6 +39,15 @@ extra_body = [
     { pointer = "/frequency_penalty", value = 1.42 },
 ]
 
+[models."gpt-oss-20b-fireworks"]
+routing = ["fireworks"]
+
+[models."gpt-oss-20b-fireworks".providers.fireworks]
+type = "fireworks"
+model_name = "accounts/fireworks/models/gpt-oss-20b"
+parse_think_blocks = false
+
+
 [models."gpt-4o-mini-2024-07-18-dynamic"]
 routing = ["openai"]