From 4fbbbd2d991b37ac04b77151ef862de9649bbfec Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 28 Jun 2024 06:24:20 +0800
Subject: feat: `.file`/`--file` support URLs (#665)

---
 src/utils/command.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

(limited to 'src/utils/command.rs')

diff --git a/src/utils/command.rs b/src/utils/command.rs
index 557649a..1a7de5b 100644
--- a/src/utils/command.rs
+++ b/src/utils/command.rs
@@ -1,6 +1,8 @@
+use super::*;
+
 use std::{collections::HashMap, env, ffi::OsStr, path::Path, process::Command};
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 
 pub fn detect_os() -> String {
     let os = env::consts::OS;
@@ -91,6 +93,59 @@ pub fn run_command_with_output<T: AsRef<OsStr>>(
     Ok((status.success(), stdout.to_string(), stderr.to_string()))
 }
 
+pub fn run_loader_command(path: &str, extension: &str, loader_command: &str) -> Result<String> {
+    let cmd_args = shell_words::split(loader_command).with_context(|| {
+        anyhow!("Invalid rag document loader '{extension}': `{loader_command}`")
+    })?;
+    let mut use_stdout = true;
+    let outpath = env::temp_dir()
+        .join(format!("aichat-output-{}", sha256(path)))
+        .display()
+        .to_string();
+    let cmd_args: Vec<_> = cmd_args
+        .into_iter()
+        .map(|mut v| {
+            if v.contains("$1") {
+                v = v.replace("$1", path);
+            }
+            if v.contains("$2") {
+                use_stdout = false;
+                v = v.replace("$2", &outpath);
+            }
+            v
+        })
+        .collect();
+    let cmd_eval = shell_words::join(&cmd_args);
+    debug!("run `{cmd_eval}`");
+    let (cmd, args) = cmd_args.split_at(1);
+    let cmd = &cmd[0];
+    if use_stdout {
+        let (success, stdout, stderr) =
+            run_command_with_output(cmd, args, None).with_context(|| {
+                format!("Unable to run `{cmd_eval}`, Perhaps '{cmd}' is not installed?")
+            })?;
+        if !success {
+            let err = if !stderr.is_empty() {
+                stderr
+            } else {
+                format!("The command `{cmd_eval}` exited with non-zero.")
+            };
+            bail!("{err}")
+        }
+        Ok(stdout)
+    } else {
+        let status = run_command(cmd, args, None).with_context(|| {
+            format!("Unable to run `{cmd_eval}`, Perhaps '{cmd}' is not installed?")
+        })?;
+        if status != 0 {
+            bail!("The command `{cmd_eval}` exited with non-zero.")
+        }
+        let contents = std::fs::read_to_string(&outpath)
+            .context("Failed to read file generated by the loader")?;
+        Ok(contents)
+    }
+}
+
 pub fn edit_file(editor: &str, path: &Path) -> Result<()> {
     let mut child = Command::new(editor).arg(path).spawn()?;
     child.wait()?;
-- 
cgit v1.2.3