feat: added the ability to read from stdin

cloudbridgeuy · Apr 12, 2023 · 361162b · 361162b
1 parent 2b5636c
commit 361162b
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 1 deletion.
diff --git a/crates/b/src/chats.rs b/crates/b/src/chats.rs
@@ -1,4 +1,5 @@
 use std::error::Error;
+use std::io::Read;
 
 use async_trait::async_trait;
 use serde_either::SingleOrVec;
@@ -34,10 +35,29 @@ impl ChatsCreateCommand {
                     .expect("No API Key provided")
                     .to_string();
                 let mut api = ChatsApi::new(api_key)?;
+
                 api.messages = vec![ChatMessage {
                     content: prompt.to_owned().join(" "),
                     role: "user".to_owned(),
                 }];
+
+                let mut stdin = Vec::new();
+                // Read from stdin if it's not a tty and don't forget to unlock `stdin`
+                {
+                    let mut stdin_lock = std::io::stdin().lock();
+                    stdin_lock.read_to_end(&mut stdin)?;
+                }
+
+                if !stdin.is_empty() {
+                    api.messages.insert(
+                        0,
+                        ChatMessage {
+                            content: String::from_utf8_lossy(&stdin).to_string(),
+                            role: "user".to_owned(),
+                        },
+                    );
+                }
+
                 api.model = model.to_owned();
                 api.max_tokens = max_tokens.to_owned();
                 api.n = *n;

diff --git a/crates/b/src/completions.rs b/crates/b/src/completions.rs
@@ -1,4 +1,7 @@
 use std::error::Error;
+use std::fmt::Write;
+use std::io::Read;
+use std::string::String;
 
 use async_trait::async_trait;
 use serde_either::SingleOrVec;
@@ -38,7 +41,35 @@ impl CompletionsCreateCommand {
                     .expect("No API key provided")
                     .to_string();
                 let mut api = CompletionsApi::new(api_key)?;
-                api.prompt = Some(SingleOrVec::Vec(prompt.clone()));
+
+                let mut stdin = Vec::new();
+                // Read from stdin if it's not a tty and don't forget to unlock `stdin`
+                {
+                    let mut stdin_lock = std::io::stdin().lock();
+                    stdin_lock.read_to_end(&mut stdin)?;
+                }
+
+                if !stdin.is_empty() {
+                    if prompt.len() == 0 {
+                        api.prompt = Some(SingleOrVec::Single(
+                            String::from_utf8_lossy(&stdin).to_string(),
+                        ));
+                    } else {
+                        let mut first = String::new();
+                        write!(
+                            first,
+                            "{}\n{}",
+                            String::from_utf8_lossy(&stdin).to_string(),
+                            prompt.first().unwrap().clone(),
+                        )?;
+                        let mut clone = prompt.clone().iter().skip(1).cloned().collect::<Vec<_>>();
+                        clone.insert(0, first);
+                        api.prompt = Some(SingleOrVec::Vec(clone));
+                    }
+                } else {
+                    api.prompt = Some(SingleOrVec::Vec(prompt.clone()));
+                }
+
                 api.model = model.to_string();
                 api.max_tokens = *max_tokens;
                 api.n = *n;

diff --git a/crates/openai/src/completions.rs b/crates/openai/src/completions.rs
@@ -143,6 +143,14 @@ impl CompletionsApi {
         if let Some(_) = &self.echo {
             return Err(error::OpenAi::InvalidSuffix);
         }
+
+        // Can't run 'suffix' with multiple prompts
+        if let Some(SingleOrVec::Vec(prompts)) = &self.prompt {
+            if prompts.len() > 1 {
+                return Err(error::OpenAi::InvalidSuffix);
+            }
+        }
+
         self.suffix = Some(suffix);
 
         log::debug!("Set suffix to {:?}", &self.suffix);
@@ -311,6 +319,8 @@ impl CompletionsApi {
             }
         };
 
+        log::debug!("Response body: {}", body);
+
         let body: Completions = match serde_json::from_str(&body) {
             Ok(body) => body,
             Err(e) => {