Merge pull request #15 from dev-msp/stream-response

Stream generation responses

Author: Shuflduf

Cargo.toml

@@ -14,8 +14,9 @@ readme = "README.md"
 keywords = ["ai", "google", "gemini"]
 
 [dependencies]
+bytes = "1.10.1"
 futures = "0.3"
-reqwest = { version = "0.12.12", features = ["json", "rustls-tls"] }
+reqwest = { version = "0.12.19", features = ["stream", "json", "rustls-tls"] }
 secrecy = "0.10"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

examples/streaming.rs

@@ -0,0 +1,29 @@
+use futures::StreamExt as _;
+use gemini_rs::types::{CodeExecutionTool, Tools};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let client = gemini_rs::Client::instance();
+    let tools = vec![Tools {
+        function_declarations: None,
+        google_search: None,
+        code_execution: Some(CodeExecutionTool {}),
+    }];
+    let mut req = client.stream_generate_content("gemini-2.5-flash-preview-05-20");
+    req.message("what's the sum of the prime numbers between 1 and 100?");
+    req.tools(tools);
+
+    let stream = req.stream().await?;
+    println!("Stream started...");
+
+    stream
+        .for_each(|chunk| async move {
+            match chunk {
+                Ok(chunk) => println!("Chunk: {:?}", chunk.candidates),
+                Err(e) => eprintln!("Error in stream: {:?}", e),
+            }
+        })
+        .await;
+
+    Ok(())
+}

src/client.rs

@@ -4,17 +4,17 @@ use std::{
     sync::{Arc, LazyLock},
 };
 
-use crate::{Chat, Error, Result, chat, types};
 use futures::FutureExt as _;
 use reqwest::Method;
 use secrecy::{ExposeSecret as _, SecretString};
-use serde::Serialize;
 
-const BASE_URI: &str = "https://generativelanguage.googleapis.com";
+use crate::{Chat, Error, Result, StreamGenerateContent, chat, types};
+
+pub(crate) const BASE_URI: &str = "https://generativelanguage.googleapis.com";
 
 pub struct Route<T> {
-    client: Client,
-    kind: T,
+    pub(crate) client: Client,
+    pub(crate) kind: T,
 }
 
 impl<T> Route<T> {
@@ -38,16 +38,11 @@ impl<T: Request> IntoFuture for Route<T> {
                 .request(T::METHOD, format!("{BASE_URI}/{self}"));
 
             if let Some(body) = self.kind.body() {
-                // Debug print the request body
-                if let Ok(body_json) = serde_json::to_string_pretty(&body) {
-                    println!("Request body: {body_json}");
-                }
                 request = request.json(&body);
             };
 
             let response = request.send().await?;
             let raw_json = response.text().await?;
-            println!("Response: {raw_json}");
 
             match serde_json::from_str::<types::ApiResponse<T::Model>>(&raw_json)? {
                 types::ApiResponse::Ok(response) => Ok(response),
@@ -58,15 +53,15 @@ impl<T: Request> IntoFuture for Route<T> {
     }
 }
 
-impl<T> Deref for Route<T> {
-    type Target = T;
+impl Deref for Route<GenerateContent> {
+    type Target = GenerateContent;
 
     fn deref(&self) -> &Self::Target {
         &self.kind
     }
 }
 
-impl<T> DerefMut for Route<T> {
+impl DerefMut for Route<GenerateContent> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         &mut self.kind
     }
@@ -83,7 +78,7 @@ impl<T: Request> std::fmt::Display for Route<T> {
 /// Covers the 20% of use cases that [Chat] doesn't
 #[derive(Clone)]
 pub struct Client {
-    inner: Arc<ClientInner>,
+    pub(crate) inner: Arc<ClientInner>,
 }
 
 impl Deref for Client {
@@ -121,14 +116,18 @@ impl Client {
         Route::new(self, GenerateContent::new(model.into()))
     }
 
+    pub fn stream_generate_content(&self, model: &str) -> Route<StreamGenerateContent> {
+        Route::new(self, StreamGenerateContent::new(model))
+    }
+
     pub fn instance() -> Client {
         static STATIC_INSTANCE: LazyLock<Client> = LazyLock::new(Client::default);
         STATIC_INSTANCE.clone()
     }
 }
 
 pub struct GenerateContent {
-    model: Box<str>,
+    pub(crate) model: Box<str>,
     pub body: types::GenerateContent,
 }
 
@@ -186,8 +185,8 @@ impl Request for GenerateContent {
         fmt.write_str(":generateContent")
     }
 
-    fn body(self) -> Option<Self::Body> {
-        Some(self.body)
+    fn body(&self) -> Option<Self::Body> {
+        Some(self.body.clone())
     }
 }
 
@@ -241,14 +240,18 @@ impl DerefMut for Formatter<'_, '_> {
 }
 
 impl<'me, 'buffer> Formatter<'me, 'buffer> {
-    fn new(formatter: &'me mut std::fmt::Formatter<'buffer>) -> Self {
+    pub(crate) fn new(formatter: &'me mut std::fmt::Formatter<'buffer>) -> Self {
         Self {
             formatter,
             is_first: true,
         }
     }
 
-    fn write_query_param(&mut self, key: &str, value: &impl std::fmt::Display) -> std::fmt::Result {
+    pub(crate) fn write_query_param(
+        &mut self,
+        key: &str,
+        value: &impl std::fmt::Display,
+    ) -> std::fmt::Result {
         if self.is_first {
             self.formatter.write_char('?')?;
             self.is_first = false;
@@ -275,7 +278,7 @@ impl<'me, 'buffer> Formatter<'me, 'buffer> {
 }
 
 pub struct ClientInner {
-    reqwest: reqwest::Client,
+    pub(crate) reqwest: reqwest::Client,
     key: SecretString,
 }
 
@@ -299,7 +302,7 @@ pub trait Request: Send + Sized + 'static {
 
     fn format_uri(&self, fmt: &mut Formatter<'_, '_>) -> std::fmt::Result;
 
-    fn body(self) -> Option<Self::Body> {
+    fn body(&self) -> Option<Self::Body> {
         None
     }
 }

src/lib.rs

@@ -48,13 +48,15 @@
 mod chat;
 mod client;
 mod error;
+mod stream;
 pub mod types;
 
 pub type Result<T> = std::result::Result<T, Error>;
 
 pub use chat::Chat;
 pub use client::Client;
 pub use error::Error;
+pub use stream::{RouteStream, StreamGenerateContent};
 
 /// Creates a new Gemini client instance using the default configuration.
 pub fn client() -> Client {