|
1 | 1 | //! Gemini streaming coverage, including the migrated example path. |
2 | 2 |
|
| 3 | +use futures::StreamExt; |
3 | 4 | use rig::client::CompletionClient; |
| 5 | +use rig::completion::{CompletionModel, GetTokenUsage}; |
4 | 6 | use rig::providers::gemini; |
5 | 7 | use rig::providers::gemini::completion::gemini_api_types::{ |
6 | | - AdditionalParameters, GenerationConfig, ThinkingConfig, ThinkingLevel, |
| 8 | + AdditionalParameters, FinishReason, GenerationConfig, ThinkingConfig, ThinkingLevel, |
7 | 9 | }; |
8 | | -use rig::streaming::StreamingPrompt; |
| 10 | +use rig::streaming::{StreamedAssistantContent, StreamingPrompt}; |
9 | 11 |
|
10 | 12 | use crate::support::{ |
11 | 13 | STREAMING_PREAMBLE, STREAMING_PROMPT, assert_nonempty_response, collect_stream_final_response, |
@@ -76,3 +78,116 @@ async fn example_streaming_prompt() { |
76 | 78 | ) |
77 | 79 | .await; |
78 | 80 | } |
| 81 | + |
| 82 | +#[tokio::test] |
| 83 | +async fn final_metadata_exposes_finish_reason_and_model_version() { |
| 84 | + super::super::support::with_gemini_cassette( |
| 85 | + "streaming/final_metadata_exposes_finish_reason_and_model_version", |
| 86 | + |client| async move { |
| 87 | + let model = client.completion_model(gemini::completion::GEMINI_2_5_FLASH); |
| 88 | + let request = model |
| 89 | + .completion_request("Reply with exactly: final metadata ok") |
| 90 | + .temperature(0.0) |
| 91 | + .build(); |
| 92 | + let mut stream = model.stream(request).await.expect("stream should start"); |
| 93 | + |
| 94 | + let mut text = String::new(); |
| 95 | + let mut final_response = None; |
| 96 | + let mut final_response_count = 0; |
| 97 | + while let Some(chunk) = stream.next().await { |
| 98 | + match chunk.expect("stream chunk should succeed") { |
| 99 | + StreamedAssistantContent::Text(delta) => text.push_str(&delta.text), |
| 100 | + StreamedAssistantContent::Final(response) => { |
| 101 | + final_response_count += 1; |
| 102 | + final_response = Some(response); |
| 103 | + } |
| 104 | + _ => {} |
| 105 | + } |
| 106 | + } |
| 107 | + |
| 108 | + assert_nonempty_response(&text); |
| 109 | + assert_eq!( |
| 110 | + final_response_count, 1, |
| 111 | + "stream should yield exactly one final response" |
| 112 | + ); |
| 113 | + let final_response = final_response.expect("stream should yield final metadata"); |
| 114 | + assert!( |
| 115 | + matches!(final_response.finish_reason, Some(FinishReason::Stop)), |
| 116 | + "expected STOP finish reason, got {:?}", |
| 117 | + final_response.finish_reason |
| 118 | + ); |
| 119 | + assert_eq!( |
| 120 | + final_response.model_version.as_deref(), |
| 121 | + Some(gemini::completion::GEMINI_2_5_FLASH), |
| 122 | + "expected resolved Gemini model version to be surfaced" |
| 123 | + ); |
| 124 | + assert!( |
| 125 | + final_response.token_usage().is_some(), |
| 126 | + "expected final response to expose token usage" |
| 127 | + ); |
| 128 | + }, |
| 129 | + ) |
| 130 | + .await; |
| 131 | +} |
| 132 | + |
| 133 | +#[tokio::test] |
| 134 | +async fn final_metadata_handles_terminal_finish_reason_chunk() { |
| 135 | + super::super::support::with_gemini_cassette( |
| 136 | + "streaming/final_metadata_handles_terminal_finish_reason_chunk", |
| 137 | + |client| async move { |
| 138 | + let model = client.completion_model(gemini::completion::GEMINI_2_5_FLASH); |
| 139 | + let request = model |
| 140 | + .completion_request("Reply with exactly: contentless final metadata ok") |
| 141 | + .temperature(0.0) |
| 142 | + .build(); |
| 143 | + let mut stream = model.stream(request).await.expect("stream should start"); |
| 144 | + |
| 145 | + let mut text = String::new(); |
| 146 | + let mut final_response = None; |
| 147 | + let mut final_response_count = 0; |
| 148 | + while let Some(chunk) = stream.next().await { |
| 149 | + match chunk.expect("stream chunk should succeed") { |
| 150 | + StreamedAssistantContent::Text(delta) => text.push_str(&delta.text), |
| 151 | + StreamedAssistantContent::Final(response) => { |
| 152 | + final_response_count += 1; |
| 153 | + final_response = Some(response); |
| 154 | + } |
| 155 | + _ => {} |
| 156 | + } |
| 157 | + } |
| 158 | + |
| 159 | + assert_eq!(text.trim(), "contentless final metadata ok"); |
| 160 | + assert_eq!( |
| 161 | + final_response_count, 1, |
| 162 | + "terminal finish chunk should yield exactly one final response" |
| 163 | + ); |
| 164 | + let final_response = final_response.expect("stream should yield final metadata"); |
| 165 | + assert!( |
| 166 | + matches!(final_response.finish_reason, Some(FinishReason::Stop)), |
| 167 | + "expected STOP finish reason from contentless terminal chunk, got {:?}", |
| 168 | + final_response.finish_reason |
| 169 | + ); |
| 170 | + assert_eq!( |
| 171 | + final_response.model_version.as_deref(), |
| 172 | + Some(gemini::completion::GEMINI_2_5_FLASH), |
| 173 | + "expected modelVersion from terminal chunks to be retained" |
| 174 | + ); |
| 175 | + let usage = final_response |
| 176 | + .token_usage() |
| 177 | + .expect("expected final response to expose token usage"); |
| 178 | + assert!( |
| 179 | + usage.input_tokens > 0, |
| 180 | + "expected positive input token usage, got {usage:?}" |
| 181 | + ); |
| 182 | + assert!( |
| 183 | + usage.output_tokens > 0, |
| 184 | + "expected positive output token usage, got {usage:?}" |
| 185 | + ); |
| 186 | + assert!( |
| 187 | + usage.total_tokens >= usage.input_tokens + usage.output_tokens, |
| 188 | + "expected total token usage to include input and output tokens, got {usage:?}" |
| 189 | + ); |
| 190 | + }, |
| 191 | + ) |
| 192 | + .await; |
| 193 | +} |
0 commit comments