refactor(stream): move streaming logic to dedicated module

* Move `StreamGenerateContent` and `RouteStream` types
* Move parsing logic for streamed responses
* Update imports in `client.rs` and `lib.rs`
* Adjust visibility of related fields/methods for inter-module access

Author: dev-msp

src/client.rs

@@ -1,24 +1,20 @@
 use std::{
     fmt::Write as _,
     ops::{Deref, DerefMut},
-    pin::Pin,
     sync::{Arc, LazyLock},
-    task::Poll,
 };
 
-use bytes::Bytes;
-use futures::{FutureExt as _, Stream};
+use futures::FutureExt as _;
 use reqwest::Method;
 use secrecy::{ExposeSecret as _, SecretString};
-use serde::ser::Error as _;
 
-use crate::{Chat, Error, Result, chat, types};
+use crate::{Chat, Error, Result, StreamGenerateContent, chat, types};
 
-const BASE_URI: &str = "https://generativelanguage.googleapis.com";
+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> {
@@ -71,46 +67,6 @@ impl DerefMut for Route<GenerateContent> {
     }
 }
 
-impl Deref for Route<StreamGenerateContent> {
-    type Target = GenerateContent;
-
-    fn deref(&self) -> &Self::Target {
-        &self.kind.0
-    }
-}
-
-impl DerefMut for Route<StreamGenerateContent> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.kind.0
-    }
-}
-
-impl Route<StreamGenerateContent> {
-    pub async fn stream(self) -> std::result::Result<RouteStream<StreamGenerateContent>, String> {
-        let url = format!("{BASE_URI}/{}", self);
-        let body = self.kind.body().clone();
-        let mut request = self
-            .client
-            .reqwest
-            .request(StreamGenerateContent::METHOD, url);
-
-        if let Some(body) = body {
-            request = request.json(&body);
-        }
-
-        let response = request.send().await.map_err(|e| e.to_string())?;
-        let stream = response.bytes_stream();
-
-        Ok(RouteStream {
-            phantom: std::marker::PhantomData,
-            stream: Box::pin(stream),
-            buffer: Vec::new(),
-            pos: 0,
-            state: ParseState::CannotAdvance,
-        })
-    }
-}
-
 impl<T: Request> std::fmt::Display for Route<T> {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut fmt = Formatter::new(fmt);
@@ -119,153 +75,10 @@ impl<T: Request> std::fmt::Display for Route<T> {
     }
 }
 
-pub struct RouteStream<T> {
-    phantom: std::marker::PhantomData<T>,
-    stream: Pin<Box<dyn Stream<Item = std::result::Result<Bytes, reqwest::Error>> + Send>>,
-    buffer: Vec<u8>,
-    pos: usize, // A cursor into the buffer.
-    state: ParseState,
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum ParseState {
-    CannotAdvance,
-    ReadingChars,
-    ReadingValue,
-    Finished,
-}
-
-#[derive(Debug)]
-enum ParseOutcome {
-    Ok(Option<types::Response>),
-    Err(serde_json::Error),
-    Eof,
-}
-
-impl RouteStream<StreamGenerateContent> {
-    fn next_char_pos(&self) -> Option<usize> {
-        self.buffer[self.pos..]
-            .iter()
-            .position(|&b| !b.is_ascii_whitespace())
-            .map(|p| self.pos + p)
-    }
-
-    fn advance_next_char(&mut self) -> Option<u8> {
-        self.pos = self.next_char_pos().unwrap_or(self.buffer.len());
-        self.buffer.get(self.pos).copied()
-    }
-
-    fn current_char(&self) -> Option<u8> {
-        self.buffer.get(self.pos).copied()
-    }
-
-    fn is_bridge_char(&self) -> bool {
-        matches!(self.current_char(), Some(b'[') | Some(b','))
-    }
-
-    fn parse_chunk(&mut self) -> ParseOutcome {
-        let mut de = serde_json::Deserializer::from_slice(&self.buffer[self.pos..])
-            .into_iter::<types::Response>();
-        match de.next() {
-            Some(Ok(value)) => {
-                self.pos += de.byte_offset();
-                ParseOutcome::Ok(Some(value))
-            }
-            Some(Err(e)) if e.is_eof() => ParseOutcome::Eof,
-            Some(Err(e)) => ParseOutcome::Err(e),
-            None => ParseOutcome::Ok(None), // No more objects to read.
-        }
-    }
-
-    fn try_parse_next(&mut self) -> Option<ParseOutcome> {
-        match self.state {
-            ParseState::CannotAdvance => None, // nothing to read
-            ParseState::ReadingChars => {
-                self.advance_next_char();
-                if self.is_bridge_char() {
-                    self.pos += 1; // Move past this '[' or ','
-                    self.state = ParseState::ReadingValue;
-                    None
-                } else if let Some(b']') = self.current_char() {
-                    // If we hit a ']', we can finish reading.
-                    self.state = ParseState::Finished;
-                    Some(ParseOutcome::Ok(None))
-                } else {
-                    None
-                }
-            }
-            ParseState::ReadingValue => {
-                self.advance_next_char();
-                // Deserialize one object from our current position.
-                let outcome = self.parse_chunk();
-                match &outcome {
-                    ParseOutcome::Ok(Some(_)) => {
-                        self.state = ParseState::ReadingChars;
-                    }
-                    ParseOutcome::Ok(None) | ParseOutcome::Err(_) => {
-                        self.state = ParseState::Finished;
-                    }
-                    ParseOutcome::Eof => {}
-                };
-                Some(outcome)
-            }
-            ParseState::Finished => None,
-        }
-    }
-}
-
-impl Stream for RouteStream<StreamGenerateContent> {
-    type Item = Result<types::Response>;
-
-    fn poll_next(
-        mut self: Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> Poll<Option<Self::Item>> {
-        loop {
-            // Housekeeping: drain the buffer if we've processed a lot.
-            if self.pos > 2048 {
-                let this_pos = self.pos;
-                self.buffer.drain(..this_pos);
-                self.pos = 0;
-            }
-
-            if let Some(ParseOutcome::Ok(Some(response))) = self.try_parse_next() {
-                return Poll::Ready(Some(Ok(response)));
-            }
-
-            // If we fell through, we need more data. Poll the underlying stream.
-            match self.stream.as_mut().poll_next(cx) {
-                Poll::Ready(Some(Ok(bytes))) => {
-                    if self.buffer.is_empty() && !bytes.is_empty() {
-                        self.state = ParseState::ReadingChars;
-                    }
-                    self.buffer.extend_from_slice(&bytes);
-                    continue; // Loop again to process new data.
-                }
-                Poll::Pending => return Poll::Pending,
-                Poll::Ready(Some(Err(e))) => {
-                    self.state = ParseState::Finished;
-                    return Poll::Ready(Some(Err(Error::Http(e))));
-                }
-                Poll::Ready(None) => {
-                    // Underlying stream ended. Check if we're in a clean state.
-                    if self.state != ParseState::Finished && self.pos < self.buffer.len() {
-                        let msg =
-                            format!("stream ended with unparsed data in state {:?}", self.state);
-                        return Poll::Ready(Some(Err(serde_json::Error::custom(msg).into())));
-                    }
-                    self.state = ParseState::Finished;
-                    return Poll::Ready(None);
-                }
-            }
-        }
-    }
-}
-
 /// 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 {
@@ -304,10 +117,7 @@ impl Client {
     }
 
     pub fn stream_generate_content(&self, model: &str) -> Route<StreamGenerateContent> {
-        Route::new(
-            self,
-            StreamGenerateContent(GenerateContent::new(model.into())),
-        )
+        Route::new(self, StreamGenerateContent::new(model))
     }
 
     pub fn instance() -> Client {
@@ -317,7 +127,7 @@ impl Client {
 }
 
 pub struct GenerateContent {
-    model: Box<str>,
+    pub(crate) model: Box<str>,
     pub body: types::GenerateContent,
 }
 
@@ -380,34 +190,6 @@ impl Request for GenerateContent {
     }
 }
 
-pub struct StreamGenerateContent(GenerateContent);
-
-impl Request for StreamGenerateContent {
-    type Model = types::Response;
-    type Body = types::GenerateContent;
-
-    const METHOD: Method = Method::POST;
-
-    fn format_uri(&self, fmt: &mut Formatter<'_, '_>) -> std::fmt::Result {
-        fmt.write_str("v1beta/")?;
-        fmt.write_str("models/")?;
-        fmt.write_str(&self.0.model)?;
-        fmt.write_str(":streamGenerateContent")
-    }
-
-    fn body(&self) -> Option<Self::Body> {
-        Some(self.0.body.clone())
-    }
-}
-
-impl std::fmt::Display for StreamGenerateContent {
-    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let mut fmt = Formatter::new(fmt);
-        self.format_uri(&mut fmt)?;
-        fmt.write_query_param("key", &self.0.model)
-    }
-}
-
 #[derive(Default)]
 pub struct Models {
     page_size: Option<usize>,
@@ -458,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;
@@ -492,7 +278,7 @@ impl<'me, 'buffer> Formatter<'me, 'buffer> {
 }
 
 pub struct ClientInner {
-    reqwest: reqwest::Client,
+    pub(crate) reqwest: reqwest::Client,
     key: SecretString,
 }
 

src/lib.rs

@@ -48,14 +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 client::{RouteStream, StreamGenerateContent};
 pub use error::Error;
+pub use stream::{RouteStream, StreamGenerateContent};
 
 /// Creates a new Gemini client instance using the default configuration.
 pub fn client() -> Client {