Implement chat and tests

Author: rikhuijzer

README.md

@@ -4,6 +4,38 @@
 
 ## Examples
 
+### Chat in Bash
+
+We can chat straight from the command line.
+For example, via the DeepInfra API:
+
+```sh
+$ DEEPINFRA_KEY="$(cat /path/to/key)"; echo "hi there" | ata chat
+```
+
+This defaults to the `meta-llama/Llama-3.3-70B-Instruct` model.
+We can also create a Bash script to provide some default settings to the chat.
+For example, create a file called `chat.sh` with the following content:
+
+```bash
+#!/usr/bin/env bash
+
+# Exit on (pipe) errors.
+set -euo pipefail
+
+export OPENAI_KEY="$(cat /path/to/key)"
+
+ata chat --model="gpt-4o"
+```
+
+and add it to your PATH.
+Now, we can use it like this:
+
+```sh
+$ echo "This is a test. Respond with 'hello'." | ata chat
+hello
+```
+
 ### Text to Speech in Bash
 
 We can read a file out loud from the command line.

src/chat.rs

@@ -0,0 +1,60 @@
+use std::fs::File;
+use std::io::Write;
+use transformrs::Message;
+use transformrs::Provider;
+
+#[derive(clap::Parser)]
+pub(crate) struct ChatArgs {
+    /// Model to use (optional)
+    #[arg(long)]
+    model: Option<String>,
+
+    /// Output file (optional)
+    #[arg(long, short = 'o')]
+    output: Option<String>,
+
+    /// Stream output
+    #[arg(long, default_value_t = true)]
+    stream: bool,
+
+    /// Raw JSON output
+    #[arg(long)]
+    raw_json: bool,
+
+    /// Language code (optional)
+    #[arg(long)]
+    language_code: Option<String>,
+}
+
+fn default_model(provider: &Provider) -> String {
+    match provider {
+        Provider::Google => "models/gemini-1.5-flash",
+        Provider::OpenAI => "gpt-4o-mini",
+        _ => "meta-llama/Llama-3.3-70B-Instruct",
+    }
+    .to_string()
+}
+
+pub(crate) async fn chat(args: &ChatArgs, key: &transformrs::Key, input: &str) {
+    let provider = key.provider.clone();
+    let model = args
+        .model
+        .clone()
+        .unwrap_or_else(|| default_model(&provider));
+    let messages = vec![Message::from_str("user", input)];
+    let resp = transformrs::chat::chat_completion(&provider, key, &model, &messages)
+        .await
+        .expect("Chat completion failed");
+    if args.raw_json {
+        let json = resp.raw_value();
+        println!("{}", json.unwrap());
+    }
+    let resp = resp.structured().expect("Could not parse response");
+    let content = resp.choices[0].message.content.clone();
+    if let Some(output) = args.output.clone() {
+        let mut file = File::create(output).unwrap();
+        file.write_all(content.to_string().as_bytes()).unwrap();
+    } else {
+        println!("{}", content);
+    }
+}

src/main.rs

@@ -1,12 +1,19 @@
+mod chat;
 mod tts;
 
+use chat::ChatArgs;
 use clap::Parser;
 use std::io::Read;
 use transformrs::Key;
 use tts::TextToSpeechArgs;
 
 #[derive(clap::Subcommand)]
 enum Commands {
+    /// OpenAI-compatible chat.
+    ///
+    /// Takes text input from stdin and chats with an AI model.
+    #[command()]
+    Chat(ChatArgs),
     /// Convert text to speech
     ///
     /// Takes text input from stdin and converts it to speech using text-to-speech models.
@@ -50,6 +57,9 @@ async fn main() {
     let key = find_single_key(keys);
 
     match args.command {
+        Commands::Chat(args) => {
+            chat::chat(&args, &key, &input).await;
+        }
         Commands::Tts(args) => {
             tts::tts(&args, &key, &input).await;
         }

tests/chat.rs

@@ -0,0 +1,43 @@
+mod common;
+
+use common::ata;
+use common::load_key;
+use predicates::prelude::*;
+use transformrs::Provider;
+
+fn canonicalize_response(text: &str) -> String {
+    text.to_lowercase()
+        .trim()
+        .trim_end_matches('.')
+        .trim_end_matches('!')
+        .to_string()
+}
+
+#[test]
+fn unexpected_argument() -> Result<(), Box<dyn std::error::Error>> {
+    let mut cmd = ata();
+    cmd.arg("foobar");
+    cmd.assert()
+        .failure()
+        .stderr(predicate::str::contains("unrecognized subcommand"));
+
+    Ok(())
+}
+
+#[test]
+fn tts_no_args() -> Result<(), Box<dyn std::error::Error>> {
+    let dir = tempfile::tempdir().unwrap();
+    let mut cmd = ata();
+    let key = load_key(&Provider::DeepInfra);
+    let cmd = cmd
+        .arg("chat")
+        .env("DEEPINFRA_KEY", key)
+        .write_stdin("This is a test. Respond with 'hello'.")
+        .current_dir(&dir);
+    let output = cmd.assert().success().get_output().stdout.clone();
+
+    let text = String::from_utf8(output.clone()).unwrap();
+    let content = canonicalize_response(&text);
+    assert_eq!(content, "hello");
+    Ok(())
+}

tests/common/mod.rs

@@ -0,0 +1,28 @@
+use assert_cmd::Command;
+use std::io::BufRead;
+use transformrs::Provider;
+
+pub fn ata() -> Command {
+    Command::cargo_bin("ata").unwrap()
+}
+
+#[allow(dead_code)]
+/// Load a key from the local .env file.
+///
+/// This is used for testing only. Expects the .env file to contain keys for providers in the following format:
+///
+/// ```
+/// DEEPINFRA_KEY="<KEY>"
+/// OPENAI_KEY="<KEY>"
+/// ```
+pub fn load_key(provider: &Provider) -> String {
+    fn finder(line: &Result<String, std::io::Error>, provider: &Provider) -> bool {
+        line.as_ref().unwrap().starts_with(&provider.key_name())
+    }
+    let path = std::path::Path::new("test.env");
+    let file = std::fs::File::open(path).expect("Failed to open .env file");
+    let reader = std::io::BufReader::new(file);
+    let mut lines = reader.lines();
+    let key = lines.find(|line| finder(line, provider)).unwrap().unwrap();
+    key.split("=").nth(1).unwrap().to_string()
+}

tests/cli.rs tests/tts.rstests/cli.rs renamed to tests/tts.rs

@@ -1,12 +1,10 @@
-use assert_cmd::Command;
+mod common;
+
+use common::ata;
+use common::load_key;
 use predicates::prelude::*;
-use std::io::BufRead;
 use transformrs::Provider;
 
-fn ata() -> Command {
-    Command::cargo_bin("ata").unwrap()
-}
-
 #[test]
 fn unexpected_argument() -> Result<(), Box<dyn std::error::Error>> {
     let mut cmd = ata();
@@ -29,26 +27,6 @@ fn help() -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
-/// Load a key from the local .env file.
-///
-/// This is used for testing only. Expects the .env file to contain keys for providers in the following format:
-///
-/// ```
-/// DEEPINFRA_KEY="<KEY>"
-/// OPENAI_KEY="<KEY>"
-/// ```
-fn load_key(provider: &Provider) -> String {
-    fn finder(line: &Result<String, std::io::Error>, provider: &Provider) -> bool {
-        line.as_ref().unwrap().starts_with(&provider.key_name())
-    }
-    let path = std::path::Path::new("test.env");
-    let file = std::fs::File::open(path).expect("Failed to open .env file");
-    let reader = std::io::BufReader::new(file);
-    let mut lines = reader.lines();
-    let key = lines.find(|line| finder(line, provider)).unwrap().unwrap();
-    key.split("=").nth(1).unwrap().to_string()
-}
-
 #[test]
 fn tts_no_args() -> Result<(), Box<dyn std::error::Error>> {
     let dir = tempfile::tempdir().unwrap();