fix: no more multiple finish reasons in stream

lib/llm/src/preprocessor.rs

@@ -764,7 +764,10 @@ impl OpenAIPreprocessor {
         let jail = JailedStream::builder()
             .tool_call_parser(tool_call_parser)
             .build();
-        jail.apply(stream)
+        let jailed_stream = jail.apply(stream);
+
+        // Post-process to set finish reason to ToolCalls for the last chunk if any tool calls were emitted
+        JailedStream::fix_finish_reason(jailed_stream)
     }
 
     // Motivation: Each transformation on the stream should be a separate step to allow for more flexibility

lib/llm/src/protocols/openai/chat_completions/jail.rs

@@ -13,6 +13,7 @@ use dynamo_parsers::tool_calling::{
 };
 use dynamo_runtime::protocols::annotated::Annotated;
 use futures::{Stream, StreamExt};
+use std::collections::HashMap;
 
 use crate::utils::{MarkerMatcher, MatchResult};
 
@@ -509,7 +510,7 @@ impl JailedStream {
                                 last_annotated_event.clone(),
                                 last_annotated_comment.clone(),
                             );
-                            let responses = self.emit_choice_emissions(tool_content_emissions, chat_response, preserved_metadata);
+                            let responses = self.emit_choice_emissions(tool_content_emissions.clone(), chat_response, preserved_metadata);
                             for emitted_response in responses {
                                 yield emitted_response;
                             }
@@ -709,14 +710,15 @@ impl JailedStream {
                 .collect();
 
             // Create choice with tool calls
-            return create_choice_stream(
+            let choice = create_choice_stream(
                 choice_index,
                 Some(Role::Assistant),
                 normal_text.as_deref().unwrap_or(""),
                 Some(tool_call_chunks),
-                Some(FinishReason::ToolCalls),
+                None,
                 None,
             );
+            return choice;
         }
 
         // No tool calls found or parsing failed, return content choice
@@ -745,6 +747,44 @@ impl JailedStream {
         }
         false
     }
+
+    /// Post-processor that sets finish_reason to ToolCalls when tool calls were emitted
+    /// This should be called after apply() to fix the finish_reason for tool call chunks
+    pub fn fix_finish_reason<S>(
+        input_stream: S,
+    ) -> impl Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send
+    where
+        S: Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send + 'static,
+    {
+        stream! {
+            tokio::pin!(input_stream);
+            let mut has_tool_calls_per_choice: HashMap<u32, bool> = HashMap::new();
+
+            while let Some(mut response) = input_stream.next().await {
+                // Track if any choice emitted tool calls
+                if let Some(ref data) = response.data {
+                    for choice in &data.choices {
+                        if choice.delta.tool_calls.is_some() {
+                            has_tool_calls_per_choice.insert(choice.index, true);
+                        }
+                    }
+                }
+
+                // If this chunk has finish_reason and the choice had tool calls, override to ToolCalls
+                if let Some(ref mut data) = response.data {
+                    for choice in &mut data.choices {
+                        if choice.finish_reason.is_some()
+                            && has_tool_calls_per_choice.get(&choice.index).copied().unwrap_or(false)
+                        {
+                            choice.finish_reason = Some(FinishReason::ToolCalls);
+                        }
+                    }
+                }
+
+                yield response;
+            }
+        }
+    }
 }
 
 /// Builder for configuring a JailedStream