Only send `Stop` event at end of google completion request (#29885)

Michael Sloan created 7 months ago

I don't think this makes much of a difference in current use, but this
more closely matches other providers and cleans up the "Response"
section of eval markdown output

Release Notes:

- N/A

Change summary

crates/language_models/src/provider/google.rs | 17 +++++++++++------
1 file changed, 11 insertions(+), 6 deletions(-)

Detailed changes

crates/language_models/src/provider/google.rs

@@ -507,12 +507,18 @@ impl GoogleEventMapper {
         events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
     ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
     {
-        events.flat_map(move |event| {
-            futures::stream::iter(match event {
-                Ok(event) => self.map_event(event),
-                Err(error) => vec![Err(LanguageModelCompletionError::Other(anyhow!(error)))],
+        events
+            .map(Some)
+            .chain(futures::stream::once(async { None }))
+            .flat_map(move |event| {
+                futures::stream::iter(match event {
+                    Some(Ok(event)) => self.map_event(event),
+                    Some(Err(error)) => {
+                        vec![Err(LanguageModelCompletionError::Other(anyhow!(error)))]
+                    }
+                    None => vec![Ok(LanguageModelCompletionEvent::Stop(self.stop_reason))],
+                })
             })
-        })
     }
 
     pub fn map_event(
@@ -578,7 +584,6 @@ impl GoogleEventMapper {
         if wants_to_use_tool {
             self.stop_reason = StopReason::ToolUse;
         }
-        events.push(Ok(LanguageModelCompletionEvent::Stop(self.stop_reason)));
         events
     }
 }