stream_rs/accumulators/
gemini.rs

1//! Accumulator for Google Gemini `streamGenerateContent` responses.
2//!
3//! Gemini streams a generation as a sequence of `GenerateContentResponse`
4//! objects. Each carries a `candidates` array; every candidate has an `index`
5//! and a `content.parts` list whose entries are either a `text` fragment or a
6//! `functionCall` (`{ name, args }`). Successive responses extend the same
7//! candidate, so a streamed answer arrives as many small `text` parts that must
8//! be concatenated, optionally followed by a `finishReason`.
9//!
10//! This accumulator folds those fragments into the final per-candidate text and
11//! function calls. Like the other accumulators it is JSON-library agnostic:
12//! parse each streamed object however you like and call the typed methods below.
13//!
14//! # Example
15//!
16//! ```
17//! use stream_rs::accumulators::gemini::GeminiAccumulator;
18//!
19//! let mut acc = GeminiAccumulator::new();
20//! acc.push_text(0, "Hel");
21//! acc.push_text(0, "lo");
22//! acc.set_finish_reason(0, "STOP");
23//! assert_eq!(acc.candidate(0).unwrap().text, "Hello");
24//! assert_eq!(acc.candidate(0).unwrap().finish_reason.as_deref(), Some("STOP"));
25//! ```
26//!
27//! Function calls are recorded in arrival order. Gemini delivers a function
28//! call's `args` as one complete JSON object (not fragmented like `OpenAI`), so
29//! each [`push_function_call`](GeminiAccumulator::push_function_call) appends a
30//! whole call:
31//!
32//! ```
33//! use stream_rs::accumulators::gemini::GeminiAccumulator;
34//!
35//! let mut acc = GeminiAccumulator::new();
36//! acc.push_function_call(0, "get_weather", r#"{"city":"Paris"}"#);
37//! let call = &acc.candidate(0).unwrap().function_calls[0];
38//! assert_eq!(call.name, "get_weather");
39//! assert_eq!(call.args, r#"{"city":"Paris"}"#);
40//! ```
41
42use alloc::borrow::ToOwned;
43use alloc::collections::BTreeMap;
44use alloc::string::String;
45use alloc::vec::Vec;
46
47/// A function call emitted by the model.
48#[derive(Debug, Clone, Default, PartialEq, Eq)]
49pub struct FunctionCall {
50    /// The function name.
51    pub name: String,
52    /// The function arguments as a JSON object string, exactly as received.
53    pub args: String,
54}
55
56/// The assembled state of a single `candidates[index]`.
57#[derive(Debug, Clone, Default, PartialEq, Eq)]
58pub struct Candidate {
59    /// The concatenated text from every `parts[].text` fragment.
60    pub text: String,
61    /// Function calls in arrival order.
62    pub function_calls: Vec<FunctionCall>,
63    /// The `finishReason`, set when the final chunk for this candidate arrives.
64    pub finish_reason: Option<String>,
65}
66
67/// Accumulates Gemini `streamGenerateContent` deltas into final candidates.
68///
69/// Candidates are stored sparsely in a [`BTreeMap`] keyed by their
70/// provider-supplied `index`, so a large or non-contiguous index never forces a
71/// dense allocation up to that index.
72#[derive(Debug, Default)]
73pub struct GeminiAccumulator {
74    candidates: BTreeMap<usize, Candidate>,
75}
76
77impl GeminiAccumulator {
78    /// Create an empty accumulator.
79    #[must_use]
80    pub fn new() -> Self {
81        Self::default()
82    }
83
84    /// Append a `parts[].text` fragment for `candidates[index]`.
85    pub fn push_text(&mut self, index: usize, fragment: &str) {
86        self.candidate_mut(index).text.push_str(fragment);
87    }
88
89    /// Record a `parts[].functionCall` for `candidates[index]`.
90    ///
91    /// `args` is the function-call arguments serialized as a JSON object
92    /// string; pass it through verbatim from your JSON parser.
93    pub fn push_function_call(&mut self, index: usize, name: &str, args: &str) {
94        self.candidate_mut(index).function_calls.push(FunctionCall {
95            name: name.to_owned(),
96            args: args.to_owned(),
97        });
98    }
99
100    /// Set the `finishReason` for `candidates[index]`.
101    pub fn set_finish_reason(&mut self, index: usize, reason: &str) {
102        self.candidate_mut(index).finish_reason = Some(reason.to_owned());
103    }
104
105    /// Borrow the assembled `candidates[index]`, if it exists.
106    #[must_use]
107    pub fn candidate(&self, index: usize) -> Option<&Candidate> {
108        self.candidates.get(&index)
109    }
110
111    /// All assembled candidates in index order, skipping gaps never seen.
112    pub fn candidates(&self) -> impl Iterator<Item = (usize, &Candidate)> {
113        self.candidates.iter().map(|(&i, c)| (i, c))
114    }
115
116    fn candidate_mut(&mut self, index: usize) -> &mut Candidate {
117        self.candidates.entry(index).or_default()
118    }
119}
stream_rs/accumulators/gemini.rs

stream_rs/accumulators/
gemini.rs
