diff options
| author | sigoden <sigoden@gmail.com> | 2023-03-03 09:15:58 +0800 |
|---|---|---|
| committer | sigoden <sigoden@gmail.com> | 2023-03-03 09:15:58 +0800 |
| commit | 9e8a5481cf9aa4dfd75d5c55c003b3813ccbca1e (patch) | |
| tree | 33dde22e55191cd5c1d91fabea1cea307f65efab /src | |
| parent | b516080a661d9e6cba5cd47fa79dc62ac086f2e5 (diff) | |
| download | aichat-9e8a5481cf9aa4dfd75d5c55c003b3813ccbca1e.tar.gz | |
feat: support save chat messages to file
Diffstat (limited to 'src')
| -rw-r--r-- | src/config.rs | 7 | ||||
| -rw-r--r-- | src/main.rs | 28 |
2 files changed, 33 insertions, 2 deletions
diff --git a/src/config.rs b/src/config.rs index c42b11b..c99a6ef 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,7 @@ -use std::{fs::read_to_string, path::Path}; +use std::{ + fs::read_to_string, + path::{Path, PathBuf}, +}; use anyhow::{anyhow, Result}; use serde::Deserialize; @@ -9,6 +12,8 @@ pub struct Config { pub api_key: String, /// What sampling temperature to use, between 0 and 2 pub temperature: Option<f64>, + /// Specify a file path to save chat messages to + pub save_path: Option<PathBuf>, /// Set proxy pub proxy: Option<String>, /// Used only for debugging diff --git a/src/main.rs b/src/main.rs index f3db0f4..18be063 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod config; +use std::fs::{File, OpenOptions}; use std::io::{stdout, Write}; use std::path::Path; use std::path::PathBuf; @@ -149,11 +150,27 @@ fn run_repl( let mut trigged_ctrlc = false; let mut output = String::new(); let mut role: Option<Role> = None; + let mut save_file: Option<File> = if let Some(path) = &config.save_path { + let file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .map_err(|err| { + anyhow!( + "Failed to create/append save_file at {}, {err}", + path.display() + ) + })?; + Some(file) + } else { + None + }; let handle_line = |line: String, line_editor: &mut Reedline, trigged_ctrlc: &mut bool, role: &mut Option<Role>, - output: &mut String| + output: &mut String, + save_file: &mut Option<File>| -> Result<bool> { if line.starts_with('.') { let (name, args) = match line.split_once(' ') { @@ -228,6 +245,13 @@ fn run_repl( } } }); + if !output.is_empty() { + if let Some(file) = save_file.as_mut() { + let _ = file.write_all( + format!("AICHAT: {input}\n\n--------\n{output}\n--------\n\n").as_bytes(), + ); + } + } } Ok(false) }; @@ -238,6 +262,7 @@ fn run_repl( &mut trigged_ctrlc, &mut role, &mut output, + &mut save_file, )?; } loop { @@ -250,6 +275,7 @@ fn run_repl( &mut trigged_ctrlc, &mut role, &mut output, + &mut save_file, )?; if quit { break; |
