Skip to content

Commit

Permalink
feat: support .edit rag-docs (#964)
Browse files Browse the repository at this point in the history
  • Loading branch information
sigoden authored Nov 3, 2024
1 parent f5cdd0f commit 6b77890
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 8 deletions.
51 changes: 50 additions & 1 deletion src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1222,12 +1222,61 @@ impl Config {
Ok(())
}

pub async fn edit_rag_docs(config: &GlobalConfig, abort_signal: AbortSignal) -> Result<()> {
let mut rag = match config.read().rag.clone() {
Some(v) => v.as_ref().clone(),
None => bail!("No RAG"),
};

let document_paths = rag.document_paths();
let temp_file = temp_file(&format!("-rag-{}", rag.name()), ".txt");
tokio::fs::write(&temp_file, &document_paths.join("\n"))
.await
.with_context(|| {
format!(
"Failed to write current document paths to '{}'",
temp_file.display()
)
})?;
let editor = config.read().editor()?;
edit_file(&editor, &temp_file)?;
let new_document_paths =
tokio::fs::read_to_string(&temp_file)
.await
.with_context(|| {
format!(
"Failed to read new document paths from '{}'",
temp_file.display()
)
})?;
let new_document_paths = new_document_paths
.split('\n')
.filter_map(|v| {
let v = v.trim();
if v.is_empty() {
None
} else {
Some(v.to_string())
}
})
.collect::<Vec<_>>();
if new_document_paths.is_empty() || new_document_paths == document_paths {
bail!("No changes")
}
rag.refresh_document_paths(&new_document_paths, config, abort_signal)
.await?;
config.write().rag = Some(Arc::new(rag));
Ok(())
}

pub async fn rebuild_rag(config: &GlobalConfig, abort_signal: AbortSignal) -> Result<()> {
let mut rag = match config.read().rag.clone() {
Some(v) => v.as_ref().clone(),
None => bail!("No RAG"),
};
rag.rebuild(config, abort_signal).await?;
let document_paths = rag.document_paths().to_vec();
rag.refresh_document_paths(&document_paths, config, abort_signal)
.await?;
config.write().rag = Some(Arc::new(rag));
Ok(())
}
Expand Down
16 changes: 11 additions & 5 deletions src/rag/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -128,17 +128,23 @@ impl Rag {
Ok(rag)
}

pub async fn rebuild(
pub fn document_paths(&self) -> &[String] {
&self.data.document_paths
}

pub async fn refresh_document_paths<T>(
&mut self,
document_paths: &[T],
config: &GlobalConfig,
abort_signal: AbortSignal,
) -> Result<()> {
debug!("rebuild rag: {}", self.name);
) -> Result<()>
where
T: AsRef<str>,
{
let loaders = config.read().document_loaders.clone();
let spinner = create_spinner("Starting").await;
let paths = self.data.document_paths.clone();
tokio::select! {
ret = self.sync_documents(loaders, &paths, Some(spinner.clone())) => {
ret = self.sync_documents(loaders, document_paths, Some(spinner.clone())) => {
spinner.stop();
ret?;
}
Expand Down
12 changes: 10 additions & 2 deletions src/repl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ lazy_static::lazy_static! {
const MENU_NAME: &str = "completion_menu";

lazy_static::lazy_static! {
static ref REPL_COMMANDS: [ReplCommand; 34] = [
static ref REPL_COMMANDS: [ReplCommand; 35] = [
ReplCommand::new(".help", "Show this help message", AssertState::pass()),
ReplCommand::new(".info", "View system info", AssertState::pass()),
ReplCommand::new(".model", "Change the current LLM", AssertState::pass()),
Expand Down Expand Up @@ -105,6 +105,11 @@ lazy_static::lazy_static! {
"Init or use the RAG",
AssertState::False(StateFlags::AGENT)
),
ReplCommand::new(
".edit rag-docs",
"Edit the RAG documents",
AssertState::True(StateFlags::RAG),
),
ReplCommand::new(
".rebuild rag",
"Rebuild the RAG to sync document changes",
Expand Down Expand Up @@ -360,8 +365,11 @@ impl Repl {
Some(("session", _)) => {
self.config.write().edit_session()?;
}
Some(("rag-docs", _)) => {
Config::edit_rag_docs(&self.config, self.abort_signal.clone()).await?;
}
_ => {
println!(r#"Usage: .edit <role|session>"#)
println!(r#"Usage: .edit <role|session|rag-docs>"#)
}
}
}
Expand Down

0 comments on commit 6b77890

Please sign in to comment.