rig_core/providers/gemini/
streaming.rs

1//! The decoder of GenerateContent replies, a whole `generateContent` body or
2//! a stream of `streamGenerateContent` chunks, read as REST JSON. The REST,
3//! Vertex AI and gRPC wires share it: each part becomes a block whose
4//! provider item is the part as Gemini sent it.
5//!
6//! ```
7//! use rig_core::providers::gemini::streaming::GenerateContentDecoder;
8//!
9//! let decoder = GenerateContentDecoder::default();
10//! # let _ = decoder;
11//! ```
12
13use serde_json::{Map, Value};
14
15use super::completion::blocked_prompt_error;
16use super::completion::{map_google_finish_reason, usage_of};
17use crate::error::ProviderError;
18use crate::json_utils::Lenient;
19use crate::message::{CallId, DocumentSourceKind, Image, MediaType, MimeType, ToolName};
20use crate::operation::{Block, Completion, Finish};
21use crate::providers::internal::wire;
22use crate::wire::{
23    AdapterVerdict, Decoder, Flow, ObservationSink, Out, WireCitation, WireEvent, WireFrame,
24};
25
26/// The recognizability markers of a `streamGenerateContent` chunk: every
27/// genuine frame carries `candidates`, `usageMetadata` and/or
28/// `promptFeedback` (a blocked prompt's only chunk may carry nothing but
29/// the feedback), and the service's in-band abort carries only `error`. A
30/// frame with any of them is a chunk. A valid ID-only frame is recognized
31/// separately as metadata; other JSON is `Unknown`.
32const RECOGNIZABLE_CHUNK_KEYS: &[&str] =
33    &["candidates", "usageMetadata", "promptFeedback", "error"];
34
35/// A GenerateContent reply document in REST JSON: the whole
36/// `generateContent` body or one `streamGenerateContent` chunk.
37#[derive(Debug, Default, serde::Deserialize)]
38#[serde(transparent)]
39pub struct GenerateContentChunk(pub Map<String, Value>);
40
41/// Decode GenerateContent replies. A text or thought part continues the
42/// block of the part before it while the kind stays the same, and a part
43/// that carries only a thought signature joins that block; any other part is
44/// a block of its own. The provider's end is held until EOF because
45/// hosted-tool rounds can report intermediate finish reasons. The
46/// candidate's grounding and recitation metadata cite the text blocks when
47/// the reply ends.
48#[derive(Debug, Default)]
49pub struct GenerateContentDecoder {
50    /// The latest `finishReason`.
51    finish: Option<String>,
52    /// The latest `usageMetadata`, as Gemini sent it.
53    usage: Option<Value>,
54    model_version: Option<String>,
55    response_id: Option<String>,
56    /// The text or thought run later parts continue, and whether it is
57    /// thought.
58    open: Option<(usize, bool)>,
59    /// Whether the open run holds a signature.
60    signed: bool,
61    /// The index of the block written last.
62    last: Option<usize>,
63    /// A signature sent alone before any block, which the next block takes.
64    signature: Option<String>,
65    /// The answer text written so far, which grounding segments cite.
66    answer: grounding::AnswerText,
67    /// The text block and byte offset in it of each part of the latest
68    /// chunk; `None` for a part that is not answer text.
69    placed: Vec<Option<(usize, usize)>>,
70    /// Where the part being written put its text.
71    placement: Option<(usize, usize)>,
72    /// The latest `groundingMetadata`'s citations, by block.
73    grounding: Vec<(usize, WireCitation)>,
74    /// How many chunks carried parts: a reply over more than one counts a
75    /// grounding segment's bytes in the whole answer, not in a part.
76    chunks: usize,
77    /// The `citationMetadata` sources of every chunk, by block.
78    recitations: Vec<(usize, WireCitation)>,
79}
80
81impl<'id> Decoder<'id, Completion> for GenerateContentDecoder {
82    type Event = GenerateContentChunk;
83
84    fn classify(&self, frame: WireFrame) -> WireEvent<GenerateContentChunk> {
85        match wire::classify_marker_keyed_frame(&frame.as_str(), RECOGNIZABLE_CHUNK_KEYS) {
86            chunk @ WireEvent::Known(_) => chunk,
87            // ID-only metadata must not create an unknown-content truncation tail.
88            _ if Self::is_analysis_only(&frame) => {
89                wire::classify_untyped_line(frame.as_str().as_bytes())
90            }
91            event => event,
92        }
93    }
94
95    fn decode(
96        &mut self,
97        GenerateContentChunk(data): GenerateContentChunk,
98        mut out: Out<'id, Completion>,
99    ) -> Result<Flow, ProviderError> {
100        let mut data = Value::Object(data);
101        // The parts move into their blocks; the rest of the chunk is read
102        // in place.
103        let parts = data
104            .pointer_mut("/candidates/0/content/parts")
105            .map(Value::take);
106        // Gemini repeats both on every chunk. Record a value when it changes:
107        // every span record counts against the exporter's attribute limit.
108        let span = tracing::Span::current();
109        if let Some(id) = data.str("responseId").filter(|id| !id.is_empty())
110            && self.response_id.as_deref() != Some(id)
111        {
112            span.record("gen_ai.response.id", id);
113            self.response_id = Some(id.to_owned());
114        }
115        if let Some(model) = data.str("modelVersion").filter(|model| !model.is_empty())
116            && self.model_version.as_deref() != Some(model)
117        {
118            span.record("gen_ai.response.model", model);
119            self.model_version = Some(model.to_owned());
120        }
121        if let Some(usage) = data.get_mut("usageMetadata") {
122            self.usage = Some(usage.take());
123        }
124        if let Some(error) = data.at("/error") {
125            // An in-band failure is a provider error, not a truncation.
126            let body = serde_json::json!({ "error": error }).to_string();
127            return Err(ProviderError::from_provider_body(body));
128        }
129        if let Some(blocked) = data.get("promptFeedback").and_then(blocked_prompt_error) {
130            return Err(blocked);
131        }
132        // The candidates, their content and its parts hold every block and
133        // the finish, so a wrongly typed one fails the reply.
134        let candidate = match data
135            .get("candidates")
136            .map(|candidates| (candidates, candidates.get(0)))
137        {
138            None | Some((Value::Null, _) | (Value::Array(_), None)) => return Ok(Flow::More),
139            Some((Value::Array(_), Some(candidate @ Value::Object(_)))) => candidate,
140            Some((Value::Array(_), Some(_))) => {
141                return Err(malformed("a candidate that is not an object"));
142            }
143            Some(_) => return Err(malformed("candidates that are not a list")),
144        };
145        // Last one wins: an intermediate `finishReason` is superseded by the
146        // reason the turn actually ended on. Proto3 JSON spells an enum
147        // value its schema does not know as its number.
148        match candidate.get("finishReason") {
149            Some(Value::String(name)) => self.finish = Some(name.clone()),
150            Some(Value::Number(number)) => self.finish = Some(format!("FINISH_REASON_{number}")),
151            _ => {}
152        }
153        let parts = match candidate.get("content") {
154            None | Some(Value::Null) => None,
155            Some(Value::Object(_)) => parts,
156            Some(_) => return Err(malformed("candidate content that is not an object")),
157        };
158        match parts {
159            None | Some(Value::Null) => {}
160            Some(Value::Array(parts)) => {
161                self.chunks += 1;
162                self.placed.clear();
163                for part in parts {
164                    self.part(part, &mut out)?;
165                    self.placed.push(self.placement.take());
166                }
167            }
168            Some(_) => return Err(malformed("candidate parts that are not a list")),
169        }
170        // Each chunk restates the grounding whole; recitation sources add up.
171        if let Some(metadata) = candidate.get("groundingMetadata") {
172            self.grounding =
173                grounding::grounding(metadata, &self.placed, &self.answer, self.chunks > 1);
174        }
175        if let Some(metadata) = candidate.get("citationMetadata") {
176            let recitations = grounding::recitations(metadata, &self.answer);
177            self.recitations.extend(recitations);
178        }
179        // A failure is final: nothing after it is read.
180        use crate::completion::FinishReason::{Length, Stop};
181        let reason = self.finish.as_deref().map(map_google_finish_reason);
182        match (reason, candidate.get("finishReason")) {
183            (None | Some(Stop | Length), _) | (_, None) => Ok(Flow::More),
184            _ => self.end(out),
185        }
186    }
187
188    /// Gemini ends a reply at EOF, not at its first finish reason: a
189    /// hosted-tool round can report one before more content arrives.
190    fn eof(&mut self, out: Out<'id, Completion>) -> Result<Flow, ProviderError> {
191        self.end(out)
192    }
193}
194
195/// A transport that delivers each reply message as its REST chunk, such as
196/// an SDK's reply transcoded once, hands the decoder the chunk itself.
197impl<'id> Decoder<'id, Completion, GenerateContentChunk> for GenerateContentDecoder {
198    type Event = GenerateContentChunk;
199
200    fn classify(&self, chunk: GenerateContentChunk) -> WireEvent<GenerateContentChunk> {
201        WireEvent::Known(chunk)
202    }
203
204    fn decode(
205        &mut self,
206        chunk: GenerateContentChunk,
207        out: Out<'id, Completion>,
208    ) -> Result<Flow, ProviderError> {
209        Decoder::<'id, Completion>::decode(self, chunk, out)
210    }
211
212    fn eof(&mut self, out: Out<'id, Completion>) -> Result<Flow, ProviderError> {
213        self.end(out)
214    }
215}
216
217impl GenerateContentDecoder {
218    /// End the reply on the finish reason the candidate holds. Without one
219    /// the reply did not end.
220    fn end(&mut self, mut out: Out<'_, Completion>) -> Result<Flow, ProviderError> {
221        let Some(reason) = self.finish.take() else {
222            return Err(ProviderError::Truncated);
223        };
224        self.close(&mut out)?;
225        let citations = std::mem::take(&mut self.grounding)
226            .into_iter()
227            .chain(std::mem::take(&mut self.recitations));
228        for (index, citation) in citations {
229            out.cite(index, citation);
230        }
231        let usage = self.usage.as_ref().map(usage_of).unwrap_or_default();
232        let model = self.model_version.take();
233        let response_id = self.response_id.take();
234        Ok(out.end(Finish {
235            usage,
236            reason: Some(map_google_finish_reason(&reason)),
237            response_id,
238            model,
239            ..Finish::default()
240        }))
241    }
242
243    /// Finish the open text or thought run: Gemini states each part whole.
244    fn close(&mut self, out: &mut Out<'_, Completion>) -> Result<(), ProviderError> {
245        self.open
246            .take()
247            .map_or(Ok(()), |(index, _)| out.finish(index))
248    }
249
250    /// Write a block at a fresh index, closing the run before it. A text or
251    /// thought block stays open as the run later parts continue; any other
252    /// is whole.
253    fn open(
254        &mut self,
255        block: Block,
256        item: Value,
257        text: &str,
258        out: &mut Out<'_, Completion>,
259    ) -> Result<(), ProviderError> {
260        self.close(out)?;
261        let index = out.fresh_index();
262        let mut item = item;
263        if let (Some(signature), Some(fields)) = (self.signature.take(), item.as_object_mut()) {
264            fields
265                .entry("thoughtSignature")
266                .or_insert(Value::String(signature));
267        }
268        self.last = Some(index);
269        self.signed = item
270            .get("thoughtSignature")
271            .and_then(Value::as_str)
272            .is_some_and(|signature| !signature.is_empty());
273        if matches!(block, Block::Text) {
274            self.answer.open(index, text);
275            self.placement = Some((index, 0));
276        }
277        let run = match &block {
278            Block::Text => Some(false),
279            Block::Reasoning { .. } => Some(true),
280            _ => None,
281        };
282        out.open(index, block, item)?;
283        out.push(index, text)?;
284        self.open = run.map(|thought| (index, thought));
285        self.open.map_or_else(|| out.finish(index), |_| Ok(()))
286    }
287
288    /// Write one part of the candidate's content. A part this decoder does
289    /// not model is an opaque item that replays; one with no data, or that
290    /// is not an object, is kept as one that does not. A signature alone
291    /// joins the block before it, as Gemini attaches it to the part before,
292    /// or the next block when it comes first.
293    fn part(&mut self, part: Value, out: &mut Out<'_, Completion>) -> Result<(), ProviderError> {
294        let Value::Object(fields) = &part else {
295            return self.open(Block::Opaque { replay: false }, part, "", out);
296        };
297        let thought = part.bool("thought") == Some(true);
298        let signature = part
299            .str("thoughtSignature")
300            .filter(|signature| !signature.is_empty());
301        if let Some(text) = part.str("text") {
302            // An empty part without a signature carries nothing.
303            if text.is_empty() && signature.is_none() {
304                return Ok(());
305            }
306            // A second signed part starts a block of its own, so each
307            // signature stays with the part that carried it.
308            return match self.open {
309                Some((index, kind)) if kind == thought && !(self.signed && signature.is_some()) => {
310                    self.signed |= signature.is_some();
311                    if !thought {
312                        self.placement = self.answer.push(index, text).map(|at| (index, at));
313                    }
314                    out.push(index, text)?;
315                    out.edit(index, |item| merge_part(item, &part))
316                }
317                _ => {
318                    let block = match thought {
319                        true => Block::Reasoning { redacted: false },
320                        false => Block::Text,
321                    };
322                    let text = text.to_owned();
323                    self.open(block, part, &text, out)
324                }
325            };
326        }
327        if let Some(call) = part.obj("functionCall") {
328            let Ok(name) =
329                ToolName::new(call.get("name").and_then(Value::as_str).unwrap_or_default())
330            else {
331                tracing::warn!("Gemini sent a function call without a name; nothing can answer it");
332                return Ok(());
333            };
334            // Proto3 JSON leaves out an empty `args` Struct.
335            let args = call
336                .get("args")
337                .map_or_else(|| "{}".to_owned(), Value::to_string);
338            let id = CallId::from_wire(call.get("id").and_then(Value::as_str).unwrap_or_default());
339            return self.open(Block::Call { id, name }, part, &args, out);
340        }
341        if !thought
342            && let (Some(mime_type), Some(data)) =
343                (part.at("/inlineData/mimeType"), part.at("/inlineData/data"))
344            && let (Some(mime_type), Some(data)) = (mime_type.as_str(), data.as_str())
345            && let Some(MediaType::Image(media_type)) = MediaType::from_mime_type(mime_type)
346        {
347            let image = Image {
348                data: DocumentSourceKind::Base64(data.to_owned()),
349                media_type: Some(media_type),
350                detail: None,
351                native: None,
352            };
353            return self.open(Block::Image(image), part, "", out);
354        }
355        let bare = ["thought", "thoughtSignature", "partMetadata"];
356        let data = fields.keys().any(|key| !bare.contains(&key.as_str()));
357        if !data && let Some(signature) = signature {
358            let Some(index) = self.open.map(|(index, _)| index).or(self.last) else {
359                self.signature = Some(signature.to_owned());
360                return Ok(());
361            };
362            return out.edit(index, |item| {
363                if let Some(item) = item.as_object_mut() {
364                    item.insert("thoughtSignature".to_owned(), Value::from(signature));
365                }
366            });
367        }
368        self.open(Block::Opaque { replay: data }, part, "", out)
369    }
370
371    /// Whether `frame` carries a response id and nothing else, a repeated
372    /// key included.
373    pub(crate) fn is_analysis_only(frame: &WireFrame) -> bool {
374        #[derive(serde::Deserialize)]
375        #[serde(deny_unknown_fields)]
376        struct ResponseIdOnly {
377            #[serde(rename = "responseId")]
378            _id: String,
379        }
380        matches!(
381            wire::classify_untyped_line::<ResponseIdOnly>(frame.as_str().as_bytes()),
382            WireEvent::Known(_)
383        )
384    }
385
386    /// Project provider verdicts, response identity, and errors
387    /// before normalization. A field of another type is left out, and a
388    /// payload that is not JSON projects nothing.
389    pub(crate) fn project(payload: &[u8], sink: &mut ObservationSink<'_>) {
390        let Ok(reply) = serde_json::from_slice::<Value>(payload) else {
391            return;
392        };
393        let candidate = reply.arr("candidates").first().unwrap_or(&Value::Null);
394        let scrub = |value: Option<&str>| value.map(|value| sink.scrub(value));
395        let block = reply
396            .at("/promptFeedback/blockReason")
397            .and_then(Value::as_str);
398        let verdict = AdapterVerdict {
399            finish_reason: scrub(candidate.str("finishReason")),
400            block_reason: scrub(block),
401            detail: scrub(candidate.str("finishMessage")),
402            model: scrub(reply.str("modelVersion")),
403        };
404        let response_id = scrub(reply.str("responseId"));
405        sink.provider(verdict, response_id);
406        if let Some(error) = crate::observe::ObservedError::of(&reply) {
407            error.emit(sink);
408        }
409    }
410}
411
412fn malformed(what: &str) -> ProviderError {
413    ProviderError::Response(format!("Gemini sent {what}"))
414}
415
416/// Merge a streamed part into the part its block holds so far: the text
417/// appends, a non-empty `thoughtSignature` replaces the one held, and any
418/// other field replaces its namesake.
419fn merge_part(item: &mut Value, part: &Value) {
420    let (Value::Object(held), Value::Object(part)) = (&mut *item, part) else {
421        *item = part.clone();
422        return;
423    };
424    for (key, value) in part {
425        if key == "text"
426            && let (Some(Value::String(text)), Value::String(more)) = (held.get_mut("text"), value)
427        {
428            text.push_str(more);
429        } else if key != "thoughtSignature" || value.as_str() != Some("") {
430            held.insert(key.clone(), value.clone());
431        }
432    }
433}
434
435pub mod document;
436mod grounding;
437
438#[cfg(test)]
439mod tests;
rig_core/providers/gemini/streaming.rs

rig_core/providers/gemini/
streaming.rs
