From 9e8a5481cf9aa4dfd75d5c55c003b3813ccbca1e Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 3 Mar 2023 09:15:58 +0800
Subject: feat: support save chat messages to file

---
 src/config.rs |  7 ++++++-
 src/main.rs   | 28 +++++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 2 deletions(-)

(limited to 'src')

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;
-- 
cgit v1.2.3