feat: redactions are reverted before returning the output

Cargo.toml

@@ -17,3 +17,6 @@ tempfile = "3.13.0"
 dirs = "5.0.1"
 regex = "1.11.0"
 
+[dependencies.uuid]
+version = "1.11.0"
+features = ["v4"]

src/main.rs

@@ -4,6 +4,7 @@ mod openai;
 mod output;
 mod path;
 mod repository;
+mod redactions;
 
 use crate::args::Args;
 use crate::config::repository::ConfigRepository;
@@ -15,7 +16,7 @@ use clap::Parser;
 use config::{model::keys::ConfigKeys, service::config_service};
 use openai::service::chat::chat;
 use output::outputter;
-use regex::Regex;
+use output::message::Message;
 use repository::db::SqliteRepository;
 use std::fs::create_dir_all;
 use std::io::IsTerminal;
@@ -92,16 +93,9 @@ async fn request_response_from_ai<R: ConfigRepository>(
         None => input.clone(),
     };
 
-    let redactions = redacted_config::fetch_redactions(repo);
-    let input_with_redactions =
-        redactions
-            .iter()
-            .fold(input_with_local_context.clone(), |acc, redaction| {
-                let re = Regex::new(&format!("(?i){}", regex::escape(redaction))).unwrap();
-                re.replace_all(&acc, "<redacted>").to_string()
-            });
+    let (redacted_input, mapped_redactions) = redactions::redact::redact(repo, &input_with_local_context);
 
-    let chat_response = match chat(&open_ai_api_key.value, user_defined_system_prompt, &input_with_redactions).await {
+    let chat_response = match chat(&open_ai_api_key.value, user_defined_system_prompt, &redacted_input).await {
         Ok(response) => response,
         Err(err) => {
             println!("{:#?}", err);
@@ -112,9 +106,14 @@ async fn request_response_from_ai<R: ConfigRepository>(
     let output_messages = chat_response
         .iter()
         .map(|message| message.to_output_message())
-        .collect();
+        .collect::<Vec<Message>>();
 
-    outputter::print(output_messages);
+    let unredacted_messages = output_messages.iter().map(|message| {
+        let unredacted_message = redactions::revert::unredact(&mapped_redactions, &message.message);
+        message.copy_with_message(unredacted_message)
+    }).collect::<Vec<Message>>();
+
+    outputter::print(unredacted_messages);
     Ok(())
 }
 

src/openai/model/role.rs

@@ -1,3 +1,4 @@
+#[derive(Clone)]
 pub enum Role {
     System,
     User,

src/output/message.rs

@@ -4,3 +4,12 @@ pub struct Message {
     pub role: Role,
     pub message: String,
 }
+
+impl Message {
+    pub fn copy_with_message(&self, message: String) -> Self {
+        return Message {
+            role: self.role.clone(),
+            message
+        }
+    }
+}

src/redactions/common.rs

@@ -0,0 +1,14 @@
+use std::collections::HashMap;
+use uuid::Uuid;
+
+pub fn redaction_map(redactions: Vec<String>) -> HashMap<String, String> {
+    let mut map = HashMap::new();
+    for redaction in redactions {
+        map.insert(redaction, generate_uuid_v4());
+    }
+    map
+}
+
+fn generate_uuid_v4() -> String {
+    Uuid::new_v4().to_string()
+}

src/redactions/mod.rs

@@ -0,0 +1,3 @@
+pub(crate) mod redact;
+pub(crate) mod revert;
+pub(crate) mod common;

src/redactions/redact.rs

@@ -0,0 +1,21 @@
+use std::collections::HashMap;
+use regex::Regex;
+use crate::config::repository::ConfigRepository;
+use crate::config::service::redacted_config;
+
+use super::common;
+
+pub fn redact<R: ConfigRepository>(repo: &R, content: &str) -> (String, HashMap<String, String>) {
+    let redactions = redacted_config::fetch_redactions(repo);
+    let mapped_redactions = common::redaction_map(redactions);
+
+    let input_with_redactions =
+        mapped_redactions
+            .iter()
+            .fold(content.to_string(), |acc, (redaction, id)| {
+                let re = Regex::new(&format!("(?i){}", regex::escape(redaction))).unwrap();
+                re.replace_all(&acc, id).to_string()
+            });
+
+    (input_with_redactions.to_string(), mapped_redactions)
+}

src/redactions/revert.rs

@@ -0,0 +1,12 @@
+use std::collections::HashMap;
+use regex::Regex;
+
+pub fn unredact(mapped_redactions: &HashMap<String, String>, content: &str) -> String {
+    mapped_redactions
+        .iter()
+        .fold(content.to_string(), |acc, (redaction, id)| {
+            let re = Regex::new(&format!("(?i){}", regex::escape(id))).unwrap();
+            re.replace_all(&acc, redaction).to_string()
+        })
+}
+