use async_std::{io, process::Command}; use futures::{future::BoxFuture, stream::BoxStream, Future, Stream, StreamExt, TryStreamExt}; pub trait SampleModel { type Error; type Sample: Future>; fn get_sample(&self) -> Self::Sample; } pub struct GPTSampleModel { python_command: String, command_working_path: String, command_args: Vec, } impl SampleModel for GPTSampleModel { type Error = io::Error; type Sample = BoxFuture<'static, Result>; fn get_sample(&self) -> Self::Sample { let cmd = Command::new(&self.python_command) .current_dir(&self.command_working_path) .args(&self.command_args) .output(); Box::pin(async { Ok(String::from_utf8_lossy(&cmd.await?.stdout).to_string()) }) } } impl GPTSampleModel { pub fn new( python_command: String, command_working_path: String, command_args: Vec, ) -> GPTSampleModel { Self { python_command: python_command, command_working_path: command_working_path, command_args: command_args, } } } pub trait SampleModelExt: SampleModel { type Stream: Stream>; fn into_stream(self) -> Self::Stream; } const SAMPLE_SPLIT_WORD: &str = "<|endoftext|>"; const SAMPLE_SAMPLE_LINE: &str = "======================================== SAMPLE 1 ========================================"; impl SampleModelExt for T where Self::Sample: Send, { type Stream = BoxStream<'static, Result>; fn into_stream(self) -> Self::Stream { Box::pin( futures::stream::try_unfold(self, |this| async { Ok(Some((this.get_sample().await?, this))) }) .map_ok(|samples| { futures::stream::iter( samples .replace(SAMPLE_SAMPLE_LINE, "") .split(SAMPLE_SPLIT_WORD) .map(|elem| elem.to_owned()) .collect::>() .into_iter() .map(|elem| Ok(elem.trim().to_owned())), ) }) .try_flatten(), ) } }