Broken streaming

llm-chain-local/Cargo.toml

@@ -13,12 +13,10 @@ repository = "https://github.com/sobelio/llm-chain/"
 
 [dependencies]
 async-trait = "0.1.68"
-futures = "0.3.28"
 llm = "0.1.0-rc3"
 llm-chain = { path = "../llm-chain", version = "0.9.0", default-features = false }
 rand = "0.8.5"
 serde = { version = "1.0.160", features = ["derive"] }
 thiserror = "1.0.40"
-
-[dev-dependencies]
 tokio = { version = "1.28.0", features = ["macros", "rt"] }
+

llm-chain-local/examples/simple.rs

@@ -33,7 +33,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     let exec = LocalExecutor::new_with_options(Some(exec_opts), None)?;
     let res = exec.execute(None, &Data::Text(String::from(prompt))).await?;
+
+    while let Some(chunk) = res.receiver.lock().unwrap().recv().await {
+      print!("{chunk}")
+    }
 
-    println!("{}", res);
     Ok(())
 }

llm-chain-local/src/executor.rs

@@ -3,7 +3,7 @@ use std::env::var;
 use std::path::Path;
 
 use crate::options::{PerExecutor, PerInvocation};
-use crate::output::Output;
+use crate::output::{Output, StreamingOutput};
 use crate::LocalLlmTextSplitter;
 
 use async_trait::async_trait;
@@ -38,7 +38,7 @@ impl ExecutorError for Error {}
 impl llm_chain::traits::Executor for Executor {
     type PerInvocationOptions = PerInvocation;
     type PerExecutorOptions = PerExecutor;
-    type Output = Output;
+    type Output = StreamingOutput;
     type Error = Error;
     type Token = i32;
     type StepTokenizer<'a> = LocalLlmTokenizer<'a>;
@@ -85,23 +85,26 @@ impl llm_chain::traits::Executor for Executor {
         prompt: &Prompt,
     ) -> Result<Self::Output, Self::Error> {
         let session = &mut self.llm.start_session(Default::default());
-        let mut output = String::new();
+
+        let (send, recv) = tokio::sync::mpsc::channel::<String>(8);
+        let mut stream_out = StreamingOutput::new(recv);
         session
-            .infer::<Infallible>(
+            .infer::<tokio::sync::mpsc::error::TrySendError<String>>(
                 self.llm.as_ref(),
                 prompt.to_text().as_str(),
                 // EvaluateOutputRequest
                 &mut Default::default(),
                 &mut rand::thread_rng(),
                 |t| {
-                    output.push_str(t);
+                    stream_out.add_chunk(String::from(t));
+                    send.blocking_send(String::from(t))?;
 
                     Ok(())
                 },
             )
             .map_err(|e| Error::InnerError(Box::new(e)))?;
 
-        Ok(output.into())
+        Ok(stream_out)
     }
 
     fn tokens_used(

llm-chain-local/src/output.rs

@@ -1,6 +1,6 @@
 use async_trait::async_trait;
 use llm_chain::output;
-use std::{fmt::{Display, Formatter}, sync::{Arc, Mutex, mpsc}};
+use std::{fmt::{Display, Formatter}, sync::{Arc, Mutex}};
 
 /// Represents the output from the LLAMA model.
 #[derive(Debug, Clone)]
@@ -10,34 +10,28 @@ pub struct Output {
 
 #[derive(Clone, Debug)]
 pub struct StreamingOutput {
-    complete: bool,
     output: Arc<Mutex<Vec<String>>>,
-    receiver: Arc<Mutex<futures::channel::mpsc::Receiver<String>>>
+    pub receiver: Arc<Mutex<tokio::sync::mpsc::Receiver<String>>>
 }
 
 impl StreamingOutput {
-    fn new() -> Self {
-        let (_, receiver) = futures::channel::mpsc::channel(8);
+    pub(crate) fn new(recv: tokio::sync::mpsc::Receiver<String>) -> Self {
         Self {
-            complete: false,
             output: Arc::new(Mutex::new(vec![])),
-            receiver: Arc::new(Mutex::new(receiver))
+            receiver: Arc::new(Mutex::new(recv))
         }
     }
 
-    fn set_complete(&mut self) {
-        self.complete = true;
-    }
-
-    fn complete(&self) -> bool {
-        self.complete
+    pub(crate) fn add_chunk(&mut self, chunk: String) {
+        let mut chunks = self.output.lock().unwrap();
+        chunks.push(chunk);
     }
 }
 
 #[async_trait]
 impl output::Output for StreamingOutput {
     async fn primary_textual_output_choices(&self) -> Vec<String> {
-        todo!()
+        self.output.lock().unwrap().clone()
     }
 }