Adding streaming function to mistralrs server. (#986)

* Adding streaming function to mistralrs server.

* Adding simple_stream example

mistralrs/examples/simple_stream/main.rs

@@ -0,0 +1,55 @@
+use anyhow::Result;
+use mistralrs::{
+    IsqType, PagedAttentionMetaBuilder, RequestBuilder, Response, TextMessageRole, TextMessages,
+    TextModelBuilder,
+};
+use std::io::Write;
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let model = TextModelBuilder::new("microsoft/Phi-3.5-mini-instruct")
+        .with_isq(IsqType::Q8_0)
+        .with_logging()
+        .with_paged_attn(|| PagedAttentionMetaBuilder::default().build())?
+        .build()
+        .await?;
+
+    let messages = TextMessages::new()
+        .add_message(
+            TextMessageRole::System,
+            "You are an AI agent with a specialty in programming.",
+        )
+        .add_message(
+            TextMessageRole::User,
+            "Hello! How are you? Please write generic binary search function in Rust.",
+        );
+
+    let response = model.send_chat_request(messages).await?;
+
+    println!("{}", response.choices[0].message.content.as_ref().unwrap());
+    dbg!(
+        response.usage.avg_prompt_tok_per_sec,
+        response.usage.avg_compl_tok_per_sec
+    );
+
+    // Next example: Return some logprobs with the `RequestBuilder`, which enables higher configurability.
+    let request = RequestBuilder::new().return_logprobs(true).add_message(
+        TextMessageRole::User,
+        "Please write a mathematical equation where a few numbers are added.",
+    );
+
+    let mut stream = model.stream_chat_request(request).await?;
+
+    let stdout = std::io::stdout();
+    let lock = stdout.lock();
+    let mut buf = std::io::BufWriter::new(lock);
+    while let Some(chunk) = stream.next().await {
+        if let Response::Chunk(chunk) = chunk {
+            buf.write(chunk.choices[0].delta.content.as_bytes())?;
+        } else {
+            // Handle errors
+        }
+    }
+
+    Ok(())
+}

mistralrs/src/lib.rs

@@ -49,6 +49,45 @@
 //!     Ok(())
 //! }
 //! ```
+//!
+//! ## Streaming example
+//! ```no_run
+//! use anyhow::Result;
+//! use mistralrs::{
+//!     IsqType, PagedAttentionMetaBuilder, TextMessageRole, TextMessages, TextModelBuilder,
+//! };
+//!
+//! #[tokio::main]
+//! async fn main() -> Result<()> {
+//!     let model = TextModelBuilder::new("microsoft/Phi-3.5-mini-instruct".to_string())
+//!         .with_isq(IsqType::Q8_0)
+//!         .with_logging()
+//!         .with_paged_attn(|| PagedAttentionMetaBuilder::default().build())?
+//!         .build()
+//!         .await?;
+//!
+//!     let messages = TextMessages::new()
+//!         .add_message(
+//!             TextMessageRole::System,
+//!             "You are an AI agent with a specialty in programming.",
+//!         )
+//!         .add_message(
+//!             TextMessageRole::User,
+//!             "Hello! How are you? Please write generic binary search function in Rust.",
+//!         );
+//!
+//!     let mut stream = model.stream_chat_request(messages).await?;
+//!
+//!     while let Some(chunk) = stream.next().await {
+//!         if let Response::Chunk(chunk) = chunk{
+//!             print!("{}", chunk.choices[0].delta.content);
+//!         }
+//!         // Handle the error cases.
+//!
+//!     }
+//!     Ok(())
+//! }
+//! ```
 
 mod anymoe;
 mod diffusion_model;

mistralrs/src/model.rs

@@ -3,7 +3,7 @@ use candle_core::{Device, Result, Tensor};
 use either::Either;
 use mistralrs_core::*;
 use std::sync::Arc;
-use tokio::sync::mpsc::channel;
+use tokio::sync::mpsc::{channel, Receiver};
 
 use crate::{RequestLike, TextMessages};
 
@@ -47,11 +47,57 @@ pub struct Model {
     runner: Arc<MistralRs>,
 }
 
+pub struct Stream<'a> {
+    _server: &'a Model,
+    rx: Receiver<Response>,
+}
+
+impl<'a> Stream<'a> {
+    pub async fn next(&mut self) -> Option<Response> {
+        self.rx.recv().await
+    }
+}
+
 impl Model {
     pub fn new(runner: Arc<MistralRs>) -> Self {
         Self { runner }
     }
 
+    /// Generate with the model.
+    pub async fn stream_chat_request<R: RequestLike>(
+        &self,
+        mut request: R,
+    ) -> anyhow::Result<Stream> {
+        let (tx, rx) = channel(1);
+
+        let (tools, tool_choice) = if let Some((a, b)) = request.take_tools() {
+            (Some(a), Some(b))
+        } else {
+            (None, None)
+        };
+        let request = Request::Normal(NormalRequest {
+            messages: request.take_messages(),
+            sampling_params: request.take_sampling_params(),
+            response: tx,
+            return_logprobs: request.return_logprobs(),
+            is_streaming: true,
+            id: 0,
+            constraint: request.take_constraint(),
+            suffix: None,
+            adapters: request.take_adapters(),
+            tools,
+            tool_choice,
+            logits_processors: request.take_logits_processors(),
+            return_raw_logits: false,
+        });
+
+        self.runner.get_sender()?.send(request).await?;
+
+        let stream = Stream { _server: self, rx };
+
+        Ok(stream)
+    }
+
     /// Generate with the model.
     pub async fn send_chat_request<R: RequestLike>(
         &self,