Skip to content

Commit 6187da4

Browse files
committed
Only send Stop event at end of gemini completion request
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
1 parent bb82d9c commit 6187da4

File tree

1 file changed

+11
-6
lines changed

1 file changed

+11
-6
lines changed

crates/language_models/src/provider/google.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -507,12 +507,18 @@ impl GoogleEventMapper {
507507
events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
508508
) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
509509
{
510-
events.flat_map(move |event| {
511-
futures::stream::iter(match event {
512-
Ok(event) => self.map_event(event),
513-
Err(error) => vec![Err(LanguageModelCompletionError::Other(anyhow!(error)))],
510+
events
511+
.map(Some)
512+
.chain(futures::stream::once(async { None }))
513+
.flat_map(move |event| {
514+
futures::stream::iter(match event {
515+
Some(Ok(event)) => self.map_event(event),
516+
Some(Err(error)) => {
517+
vec![Err(LanguageModelCompletionError::Other(anyhow!(error)))]
518+
}
519+
None => vec![Ok(LanguageModelCompletionEvent::Stop(self.stop_reason))],
520+
})
514521
})
515-
})
516522
}
517523

518524
pub fn map_event(
@@ -578,7 +584,6 @@ impl GoogleEventMapper {
578584
if wants_to_use_tool {
579585
self.stop_reason = StopReason::ToolUse;
580586
}
581-
events.push(Ok(LanguageModelCompletionEvent::Stop(self.stop_reason)));
582587
events
583588
}
584589
}

0 commit comments

Comments
 (0)