use anyhow::Result;
use clap::Parser;
use ragrig::{
ChatAgentSpec, ChunkConfig, DocumentParser, DocumentParsers, DocumentType,
EmbedderSpec, EpubParserBackend, FsSessionStore, GenerationParams,
HistoryStrategy, LogHistory, PaperResult, PdfParserBackend, RagAgent, RagrigError,
ScoredChunk, SessionId, SessionStore, SummaryHistory,
Turn, TurnRole,
collect_documents, collect_documents_with_stats, download_and_ingest_url, embed_documents,
search_arxiv, search_semantic_scholar,
};
use ragrig::types::{Args, ContextSizeMode, FileHashEntry, Provider};
use ragrig::documents::{HashMetadata, get_document_file_hashes, get_changed_documents, update_file_hashes};
use ragrig::vector::{get_embeddings_file_path, remove_deleted_embeddings};
use ragrig::{parsers, store};
use rustyline::DefaultEditor;
use rustyline::error::ReadlineError;
use std::fs;
use std::io::{Write, stdout};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
struct Session {
args: Args,
agent: RagAgent,
embeddings_file_path: PathBuf,
last_results: Vec<ScoredChunk>,
last_search_results: Vec<PaperResult>,
rl: DefaultEditor,
history_path: PathBuf,
http_client: reqwest::Client,
prompt_memory: Vec<Turn>,
session_store: Box<dyn SessionStore>,
session_id: SessionId,
history_strategy: Option<Box<dyn HistoryStrategy>>,
doc_parsers: DocumentParsers,
pdf_parser: PdfParserBackend,
epub_parser: EpubParserBackend,
context_size_forced: ContextSizeMode,
}
enum Command {
#[allow(dead_code)]
Download(String),
GetPapers(String),
Help,
Scholar(String),
SearchArxiv(String),
ExtractRefs(String),
Chat(String),
Search(String),
Embed(String),
Memory(String),
Hist(String),
Parser(String),
Prompt(String),
RagQuery(String),
Unknown(String),
Exit,
}
fn filtered_parsers(pdf: &PdfParserBackend, _sloppy_pdf: bool) -> Vec<Box<dyn DocumentParser>> {
#[allow(deprecated)]
let selected_pdf = match pdf {
#[cfg(feature = "kreuzberg")]
PdfParserBackend::Kreuzberg => "kreuzberg",
PdfParserBackend::Unpdf => "unpdf",
PdfParserBackend::Sink => "pdfsink",
PdfParserBackend::Extract => "pdf-extract",
PdfParserBackend::Internal => "sloppy-pdf",
PdfParserBackend::Vision => "vision-pdf",
};
let fallback = {
#[cfg(feature = "kreuzberg")]
{ "kreuzberg" }
#[cfg(not(feature = "kreuzberg"))]
{ "sloppy-pdf" }
};
let mut list = parsers::build_parsers();
list.retain(|p| {
if p.extensions().contains(&"pdf") {
p.name() == selected_pdf || p.name() == fallback
} else {
true
}
});
list
}
async fn bootstrap(args: Args) -> Result<Session> {
let chat_params = GenerationParams {
temperature: args.temperature,
top_p: args.top_p,
max_tokens: args.max_tokens,
seed: args.seed,
};
if chat_params.temperature.is_some()
|| chat_params.top_p.is_some()
|| chat_params.max_tokens.is_some()
|| chat_params.seed.is_some()
{
print!("Params:");
if let Some(t) = chat_params.temperature {
print!(" temperature={:.2}", t);
}
if let Some(p) = chat_params.top_p {
print!(" top_p={:.2}", p);
}
if let Some(n) = chat_params.max_tokens {
print!(" max_tokens={}", n);
}
if let Some(s) = chat_params.seed {
print!(" seed={}", s);
}
println!();
}
let initial_spec = match args.provider {
Provider::Ollama => ChatAgentSpec::ollama(args.model.clone(), chat_params.clone()),
Provider::Deepseek => ChatAgentSpec::deepseek(
args.deepseek_model.clone(),
args.deepseek_api_key.clone(),
chat_params.clone(),
),
};
let chat_agent = initial_spec.build()?;
println!(
"Chat: {} ({}) | chunk_size={}, chunk_overlap={}",
chat_agent.backend_name(),
chat_agent.model_name(),
args.chunk_size,
args.chunk_overlap
);
let embedder_spec = EmbedderSpec::from_args(&args);
let embedder = embedder_spec.build()?;
println!(
"Embed: {} ({})",
embedder.backend_name(),
embedder.model_name()
);
let embeddings_file_path = get_embeddings_file_path(&args.folder);
let doc_parsers = DocumentParsers::new(filtered_parsers(&args.pdf_parser, args.sloppy_pdf));
println!(
"Parsers: {} | Active PDF: {:?} | Chunker: markdown-structural",
doc_parsers.names().join(", "),
args.pdf_parser
);
let current_file_hashes = match get_document_file_hashes(&args.folder) {
Ok(hashes) => {
println!("Found {} document files with hashes.", hashes.len());
hashes
}
Err(e) => {
eprintln!("Warning: Could not compute file hashes: {}", e);
Vec::new()
}
};
let store = store::open_store(&args.folder).await?;
let chunk_cfg = ChunkConfig { size: args.chunk_size, overlap: args.chunk_overlap };
if store.is_empty() {
println!("No existing store found. Creating new one...");
collect_documents(&*embedder, &doc_parsers, &args.folder, &chunk_cfg, &*store).await?;
} else {
println!(
"Found existing store ({} chunks). Checking for changes...",
store.len()
);
let mut stored_hashes: Vec<FileHashEntry> = Vec::new();
if embeddings_file_path.exists() {
match fs::read_to_string(&embeddings_file_path) {
Ok(json) => {
if let Ok(metadata) = serde_json::from_str::<HashMetadata>(&json) {
stored_hashes = metadata.file_hashes;
}
}
Err(e) => eprintln!("Warning: Could not read hash metadata: {}", e),
}
}
if stored_hashes.is_empty() {
println!("No hash metadata found. Regenerating all embeddings...");
for source in store.sources() {
store.delete_by_source(&source).await?;
}
collect_documents(&*embedder, &doc_parsers, &args.folder, &chunk_cfg, &*store).await?;
} else {
let changed_files = get_changed_documents(¤t_file_hashes, &stored_hashes);
if !changed_files.is_empty() {
println!("Found {} changed/new files.", changed_files.len());
remove_deleted_embeddings(&*store, ¤t_file_hashes).await?;
for (_doc_type, file_name) in &changed_files {
store.delete_by_source(file_name).await?;
}
let changed_with_types: Vec<(DocumentType, String)> = changed_files
.into_iter()
.map(|(doc_type, _)| {
let file_name = doc_type.file_name().to_string();
(doc_type, file_name)
})
.collect();
embed_documents(&*embedder, &doc_parsers, &chunk_cfg, changed_with_types, &*store)
.await?;
println!("Database updated.");
} else {
println!("No files have changed. Using existing embeddings.");
}
}
}
update_file_hashes(¤t_file_hashes, &embeddings_file_path)?;
let row_count = store.len();
if row_count == 0 {
return Err(anyhow::anyhow!(ragrig::RagrigError::NoDocumentsFound {
folder: args.folder.to_string_lossy().into_owned(),
}));
}
println!("Vector store initialized with {} total entries.", row_count);
let memory_spec = ChatAgentSpec::ollama(args.memory_model.clone(), chat_params.clone());
let memory_agent = memory_spec.build()?;
println!(
"Memory: {} ({})",
memory_agent.backend_name(),
memory_agent.model_name()
);
let mut agent_builder = RagAgent::builder()
.chat(chat_agent)
.embed(embedder)
.store(store)
.rewriter(memory_agent)
.context_tokens(args.model_ctx_tokens)
.top_k(args.top_k)
.similarity_threshold(args.similarity_threshold);
if let Some(ref path) = args.prompt_chat {
let prompt_text = fs::read_to_string(path)?;
agent_builder = agent_builder.system_prompt(prompt_text);
}
if let Some(ref path) = args.prompt_rewrite {
let rewrite_text = fs::read_to_string(path)?;
agent_builder = agent_builder.rewrite_prompt(rewrite_text);
}
let agent = agent_builder.build();
let pdf_parser = args.pdf_parser.clone();
let context_size_forced = args.context_size_forced;
let mut rl = DefaultEditor::new()?;
let history_path = args.folder.join(".ragrig_history");
if history_path.exists()
&& let Err(e) = rl.load_history(&history_path) {
eprintln!("Warning: Could not load history: {}", e);
}
println!("\nRAG System Online. Commands: /download <url> | /get <nums> | /help | exit");
println!(
"Ask questions based on your loaded documents (Arrow-Up for history, Ctrl+C to exit):"
);
let sessions_dir = args.folder.join(".ragrig").join("sessions");
let session_store: Box<dyn SessionStore> =
Box::new(FsSessionStore::new(sessions_dir)?);
let session_id = SessionId(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| format!("{}", d.as_secs()))
.unwrap_or_else(|_| "0".to_string()),
);
println!("Session: {}", session_id.0);
Ok(Session {
args,
agent,
embeddings_file_path,
last_results: Vec::new(),
last_search_results: Vec::new(),
rl,
history_path,
http_client: reqwest::Client::new(),
prompt_memory: Vec::new(),
session_store,
session_id,
history_strategy: None,
doc_parsers,
pdf_parser,
epub_parser: EpubParserBackend::Epub,
context_size_forced,
})
}
impl From<&str> for Command {
fn from(input: &str) -> Self {
let input = input.trim();
if !input.starts_with('/') {
return Command::RagQuery(input.to_string());
}
if input == "exit" || input == "quit"
|| input == "/exit" || input == "/bye"
{
return Command::Exit;
}
if input == "/help" {
return Command::Help;
}
let after = |prefix: &str| -> &str {
if input.len() > prefix.len() + 1 {
&input[prefix.len()..]
} else {
""
}
};
if input.starts_with("/download ") {
let url = strip_ansi(after("/download ")).trim().to_string();
return Command::Download(url);
}
if input.starts_with("/get ") {
return Command::GetPapers(after("/get ").trim().to_string());
}
if input.starts_with("/scholar ") {
return Command::Scholar(after("/scholar ").trim().to_string());
}
if input.eq_ignore_ascii_case("/scholar") {
return Command::Scholar(String::new());
}
if input.eq_ignore_ascii_case("/search") {
return Command::Search(String::new());
}
if input.starts_with("/search ") {
return Command::Search(after("/search ").trim().to_string());
}
if input.starts_with("/arxiv ") {
return Command::SearchArxiv(after("/arxiv ").trim().to_string());
}
if input.starts_with("/refs") {
return Command::ExtractRefs(after("/refs").trim().to_string());
}
if input.starts_with("/chat") {
return Command::Chat(after("/chat").trim().to_string());
}
if input.starts_with("/embed") {
return Command::Embed(after("/embed").trim().to_string());
}
if input.starts_with("/memory") {
return Command::Memory(after("/memory").trim().to_string());
}
if input.starts_with("/hist") {
return Command::Hist(after("/hist").trim().to_string());
}
if input.starts_with("/prompt") {
return Command::Prompt(after("/prompt").trim().to_string());
}
if input.starts_with("/parser") {
return Command::Parser(after("/parser").trim().to_string());
}
Command::Unknown(input.to_string())
}
}
impl Session {
async fn auto_save(&self) -> Result<()> {
let config = ragrig::SessionConfig {
chat_backend: self.agent.chat_agent().backend_name().to_string(),
chat_model: self.agent.chat_agent().model_name().to_string(),
embed_backend: self.agent.embedder().backend_name().to_string(),
embed_model: self.agent.embedder().model_name().to_string(),
memory_strategy: if self.agent.rewriter().is_some() {
ragrig::MemoryStrategyKind::Rewrite
} else {
ragrig::MemoryStrategyKind::Off
},
memory_backend: String::new(),
memory_model: String::new(),
top_k: self.agent.top_k(),
similarity_threshold: self.agent.similarity_threshold(),
model_ctx_tokens: self.agent.context_tokens(),
};
let data = ragrig::SessionData {
id: self.session_id.clone(),
created: std::time::UNIX_EPOCH,
updated: std::time::SystemTime::now(),
config,
turns: self.prompt_memory.clone(),
};
self.session_store.save(&data).await
}
async fn execute(&mut self, cmd: Command) -> Result<()> {
match cmd {
Command::Download(url) => self.cmd_download(&url).await,
Command::GetPapers(range) => self.cmd_get_papers(&range).await,
Command::Help => {
self.cmd_help();
Ok(())
}
Command::Scholar(q) => self.cmd_search_scholar(&q).await,
Command::SearchArxiv(q) => self.cmd_search_arxiv(&q).await,
Command::ExtractRefs(filter) => self.cmd_extract_refs(&filter).await,
Command::Chat(args_str) => self.cmd_chat(&args_str).await,
Command::Search(args) => self.cmd_search(&args).await,
Command::Embed(args_str) => self.cmd_embed(&args_str).await,
Command::Memory(args_str) => self.cmd_memory(&args_str).await,
Command::Hist(args_str) => self.cmd_hist(&args_str).await,
Command::Prompt(args_str) => self.cmd_prompt(&args_str).await,
Command::Parser(args_str) => self.cmd_parser(&args_str).await,
Command::RagQuery(q) => self.cmd_rag_query(&q).await,
Command::Unknown(cmd) => {
println!("Unknown command: '{}'", cmd);
Ok(())
}
Command::Exit => Ok(()),
}
}
async fn cmd_download(&mut self, url: &str) -> Result<()> {
if url.is_empty() {
println!("Usage: /download <url>");
return Ok(());
}
println!("Downloading and ingesting: {} ...", url);
eprintln!("[DEBUG] URL bytes: {:?}", url.as_bytes());
match download_and_ingest_url(
&*self.agent.embedder(),
&self.doc_parsers,
&self.args.folder,
&ChunkConfig { size: self.args.chunk_size, overlap: self.args.chunk_overlap },
&self.http_client,
self.agent.store(),
url,
)
.await
{
Ok(summary) => {
println!("{}", summary);
update_file_hashes(
&get_document_file_hashes(&self.args.folder).unwrap_or_default(),
&self.embeddings_file_path,
)?;
}
Err(e) => println!("Error: {}", e),
}
Ok(())
}
async fn cmd_get_papers(&mut self, range_str: &str) -> Result<()> {
if self.last_search_results.is_empty() {
println!("No search results available. Run /scholar or /arxiv first.");
return Ok(());
}
if range_str.is_empty() {
println!("Usage: /get 1,2,3-4,8");
return Ok(());
}
let indices = match parse_number_range(range_str) {
Ok(ids) => ids,
Err(e) => {
println!("Invalid range: {}", e);
return Ok(());
}
};
let mut downloaded = 0;
let mut failed = 0;
for idx in &indices {
if *idx >= self.last_search_results.len() {
println!(
" Skipping [{}]: out of range (max {})",
idx + 1,
self.last_search_results.len()
);
failed += 1;
continue;
}
let paper = &self.last_search_results[*idx];
let url = strip_ansi(&paper.best_pdf_url());
if url.is_empty() {
println!(
" [{:2}] {} — no download URL available",
idx + 1,
paper.title
);
failed += 1;
continue;
}
print!(" [{:2}] {} ... ", idx + 1, paper.title);
stdout().flush()?;
match download_and_ingest_url(
&*self.agent.embedder(),
&self.doc_parsers,
&self.args.folder,
&ChunkConfig { size: self.args.chunk_size, overlap: self.args.chunk_overlap },
&self.http_client,
self.agent.store(),
&url,
)
.await
{
Ok(_) => {
println!("done");
downloaded += 1;
update_file_hashes(
&get_document_file_hashes(&self.args.folder).unwrap_or_default(),
&self.embeddings_file_path,
)?;
}
Err(e) => {
println!("failed: {}", e);
failed += 1;
}
}
}
println!(
"Download complete: {} added, {} failed, {} skipped.",
downloaded,
failed,
indices.len().saturating_sub(downloaded + failed)
);
Ok(())
}
fn cmd_help(&self) {
println!("/download <url> — download and ingest a PDF into the document pool");
println!("/scholar <q> — search Semantic Scholar (free API key for higher limits)");
println!("/arxiv <q> — search arXiv (no API key needed, no rate limits)");
println!("/search — show / adjust vector search parameters (topk, threshold)");
println!("/get 1,2,3-4 — download papers by number from last search");
println!(
"/refs [topic] — extract references from last query results (optionally filtered by topic)"
);
println!("/chat <backend> [model] [api_key] | context <N> — hot-swap chat engine or adjust context window");
println!("/embed <backend> [model] | purge | index — hot-swap embedding backend");
println!("/memory <backend> [model] [key] | transcript | log | summary | off | purge — hot-swap memory + history diffusion");
println!("/hist [list | load <id> | delete <id>] — manage saved sessions");
println!("/prompt chat|rewrite <file> | reset — load custom system prompts");
println!(
"/parser pdf unpdf|sink|extract|internal | epub epub — hot-swap parser per format"
);
println!("exit / quit — end the session");
}
async fn cmd_search_scholar(&mut self, q: &str) -> Result<()> {
if q.is_empty() {
println!("Usage: /scholar <query>");
return Ok(());
}
println!("Searching Semantic Scholar for: {} ...", q);
match search_semantic_scholar(self.args.semantic_scholar_api_key.as_deref(), &self.http_client, q, 20).await {
Ok(papers) if papers.is_empty() => {
println!("No papers found.");
}
Ok(papers) => {
println!("\nResults:");
for (i, p) in papers.iter().enumerate() {
println!(
" [{:2}] {} — {}{}",
i + 1,
p.title,
p.format_authors(),
p.format_year()
);
let url = p.best_pdf_url();
if !url.is_empty() {
println!(" /download {}", url);
}
}
self.last_search_results = papers.clone();
println!("\nUse /download <url> to ingest any paper.");
}
Err(e) => println!("Search error: {}", e),
}
Ok(())
}
async fn cmd_search_arxiv(&mut self, q: &str) -> Result<()> {
if q.is_empty() {
println!("Usage: /arxiv <query>");
return Ok(());
}
println!("Searching arXiv for: {} ...", q);
match search_arxiv(&self.http_client, q, 20).await {
Ok(papers) if papers.is_empty() => {
println!("No papers found.");
}
Ok(papers) => {
println!("\nResults (arXiv):");
for (i, p) in papers.iter().enumerate() {
println!(
" [{:2}] {} — {}{}",
i + 1,
p.title,
p.format_authors(),
p.format_year()
);
let url = p.best_pdf_url();
if !url.is_empty() {
println!(" /download {}", url);
}
}
self.last_search_results = papers;
println!("\nUse /download <url> to ingest any paper.");
}
Err(e) => println!("arXiv search error: {}", e),
}
Ok(())
}
async fn cmd_search(&mut self, args_str: &str) -> Result<()> {
let mut parts = args_str.split_whitespace();
let sub = parts.next().unwrap_or("");
if sub.is_empty() {
println!(
"Vector search parameters:"
);
println!(" top-k: {} (change: /search topk <N>)", self.agent.top_k());
println!(" threshold: {:.3} (change: /search threshold <F>)", self.agent.similarity_threshold());
return Ok(());
}
if sub == "topk" {
match parts.next().and_then(|s| s.parse::<usize>().ok()) {
Some(n) if n > 0 => {
self.agent.set_top_k(n);
println!("Top-k set to {}.", n);
}
_ => println!("Usage: /search topk <N> (current: {})", self.agent.top_k()),
}
return Ok(());
}
if sub == "threshold" {
match parts.next().and_then(|s| s.parse::<f64>().ok()) {
Some(f) if f >= 0.0 => {
self.agent.set_similarity_threshold(f);
println!("Similarity threshold set to {:.3}.", f);
}
_ => println!(
"Usage: /search threshold <F> (current: {:.3})",
self.agent.similarity_threshold()
),
}
return Ok(());
}
println!(
"Unknown subcommand: '{}'. Use /search, /search topk <N>, or /search threshold <F>.",
sub
);
Ok(())
}
async fn cmd_extract_refs(&mut self, filter: &str) -> Result<()> {
if self.last_results.is_empty() {
println!("No previous query results. Ask a question first, then use /refs.");
return Ok(());
}
let filter_hint = if filter.is_empty() {
String::new()
} else {
format!(
" Focus specifically on references related to: \"{}\".",
filter
)
};
let mut context = String::new();
for (i, sc) in self.last_results.iter().take(5).enumerate() {
context.push_str(&format!(
"[Document {} | Source: {}]\n{}\n\n",
i + 1,
sc.chunk.source_file,
sc.chunk.text
));
}
let extract_prompt = format!(
"Extract all academic paper references (cited works with title, authors, year) from the documents below.{}\n\n\
Return ONLY a numbered list. For each reference, include:\n\
- Title of the cited paper\n\
- Authors (last name of first author + et al. if multiple)\n\
- Year\n\
- If an arXiv ID or DOI is visible, include it as a URL.\n\n\
Documents:\n{}",
filter_hint, context
);
println!("Extracting references...\n");
print!("Assistant > ");
stdout().flush()?;
let got_response = AtomicBool::new(false);
match self
.agent
.chat_agent()
.generate_stream(&extract_prompt, &|text: String| {
print!("{}", text);
let _ = stdout().flush();
got_response.store(true, Ordering::Relaxed);
})
.await
{
Ok(()) => {}
Err(e) => eprintln!("\n[ERROR] Reference extraction failed: {}", e),
}
if !got_response.load(Ordering::Relaxed) {
println!("(no references found)");
}
println!();
Ok(())
}
async fn cmd_chat(&mut self, args_str: &str) -> Result<()> {
let mut parts = args_str.split_whitespace();
let backend = parts.next().unwrap_or("");
if backend == "context" {
match parts.next().and_then(|s| s.parse::<usize>().ok()) {
Some(n) if n > 0 => {
self.agent.set_context_tokens(n);
let ctx = self.agent.context_tokens();
println!(
"Context window set to {} tokens (prompt budget ~{} chars).",
ctx,
(ctx.saturating_sub(1024)).saturating_mul(3)
);
}
_ => println!(
"Usage: /chat context <tokens> (current: {})",
self.agent.context_tokens()
),
}
return Ok(());
}
if backend == "temperature" {
match parts.next().and_then(|s| s.parse::<f64>().ok()) {
Some(t) if t >= 0.0 => {
self.rebuild_chat_with_param(|p| p.temperature = Some(t));
}
_ => println!(
"Usage: /chat temperature <F> (0.0 = deterministic)"
),
}
return Ok(());
}
if backend == "top_p" {
match parts.next().and_then(|s| s.parse::<f64>().ok()) {
Some(p) if p >= 0.0 && p <= 1.0 => {
self.rebuild_chat_with_param(|gp| gp.top_p = Some(p));
}
_ => println!(
"Usage: /chat top_p <F> (0.0–1.0)"
),
}
return Ok(());
}
if backend == "max_tokens" {
match parts.next().and_then(|s| s.parse::<usize>().ok()) {
Some(n) if n > 0 => {
self.rebuild_chat_with_param(|p| p.max_tokens = Some(n));
}
_ => println!(
"Usage: /chat max_tokens <N>"
),
}
return Ok(());
}
if backend == "seed" {
match parts.next().and_then(|s| s.parse::<u64>().ok()) {
Some(s) => {
self.rebuild_chat_with_param(|p| p.seed = Some(s));
}
_ => println!(
"Usage: /chat seed <N>"
),
}
return Ok(());
}
if backend.is_empty() {
println!(
"Chat: {} ({}) — context window: {} tokens",
self.agent.chat_agent().backend_name(),
self.agent.chat_agent().model_name(),
self.agent.context_tokens(),
);
println!("Usage: /chat <backend> [model] [api_key] | context <N> | temperature <F> | top_p <F> | max_tokens <N> | seed <N>");
println!(" backends: ollama, deepseek");
return Ok(());
}
let model = parts.next();
let api_key = parts.next();
let spec = match ChatAgentSpec::parse(backend, model, api_key, None) {
Ok(s) => s,
Err(e) => {
println!("Error: {}", e);
return Ok(());
}
};
match spec.build() {
Ok(new_agent) => {
let old_backend = self.agent.chat_agent().backend_name();
let old_model = self.agent.chat_agent().model_name().to_string();
self.agent.set_chat_agent(new_agent);
println!(
"Chat agent swapped: {} ({}) → {} ({})",
old_backend,
old_model,
self.agent.chat_agent().backend_name(),
self.agent.chat_agent().model_name()
);
}
Err(e) => {
println!("Failed to build chat agent: {}", e);
}
}
Ok(())
}
fn rebuild_chat_with_param(&mut self, f: impl FnOnce(&mut GenerationParams)) {
let current = self.agent.chat_agent();
let backend = current.backend_name();
let model = current.model_name();
let mut params = GenerationParams::default();
f(&mut params);
let spec = match ChatAgentSpec::parse(backend, Some(model), None, Some(params)) {
Ok(spec) => spec,
Err(e) => {
println!("Cannot rebuild unknown backend: {}", e);
return;
}
};
match spec.build() {
Ok(new_agent) => {
self.agent.set_chat_agent(new_agent);
let agent = self.agent.chat_agent();
println!(
"Chat params updated: {} ({})",
agent.backend_name(),
agent.model_name()
);
}
Err(e) => {
println!("Failed to rebuild chat agent: {}", e);
}
}
}
async fn cmd_embed(&mut self, args_str: &str) -> Result<()> {
let mut parts = args_str.split_whitespace();
let backend = parts.next().unwrap_or("");
if backend.is_empty() {
println!(
"Embed: {} ({}) — top‑k: {}, threshold: {}",
self.agent.embedder().backend_name(),
self.agent.embedder().model_name(),
self.agent.top_k(),
self.agent.similarity_threshold(),
);
println!(
"Usage: /embed <backend> [model] | purge | index | topk <N> | threshold <F>"
);
println!(
" backends: {}",
EmbedderSpec::available_backends().join(", ")
);
return Ok(());
}
if backend.eq_ignore_ascii_case("purge") {
let store = self.agent.store();
let count = store.len();
let sources: Vec<_> = store.sources().into_iter().collect();
for source in &sources {
store.delete_by_source(source).await?;
}
println!(
"Vector store purged ({} chunks across {} source files).",
count,
sources.len()
);
return Ok(());
}
if backend.eq_ignore_ascii_case("index") {
println!(
"Re-indexing all documents in {}...",
self.args.folder.display()
);
let chunk_cfg = ChunkConfig { size: self.args.chunk_size, overlap: self.args.chunk_overlap };
let stats = collect_documents_with_stats(&*self.agent.embedder(), &self.doc_parsers, &self.args.folder, &chunk_cfg, self.agent.store()).await?;
println!(
"Re-indexing complete. Store size: {} chunks.\n",
self.agent.store().len()
);
let ok_count = stats.iter().filter(|s| s.ok).count();
let fail_count = stats.len() - ok_count;
let total_chunks: usize = stats.iter().map(|s| s.chunks).sum();
let total_chars: usize = stats.iter().map(|s| s.chars).sum();
let total_kb: u64 = stats.iter().map(|s| s.file_size_kb).sum();
println!(
"\n{} files processed ({} ok, {} failed), {} chunks, {} chars, {} KB total.\n",
stats.len(), ok_count, fail_count, total_chunks, total_chars, total_kb
);
if !stats.is_empty() {
println!(
"{:<44} {:>6} {:>7} {:>8} {:>6}",
"File", "KB", "Chunks", "Chars", "Avg/Ch"
);
println!("{}", "─".repeat(78));
for s in &stats {
let name = if s.file_name.len() > 42 {
format!("{}…", &s.file_name[..41])
} else {
s.file_name.clone()
};
if s.ok {
println!(
"{:<44} {:>6} {:>7} {:>8} {:>6.0}",
name, s.file_size_kb, s.chunks, s.chars, s.avg_chars_per_chunk()
);
} else {
println!(
"{:<44} {:>6} {:>7} {:>8} {:>6} FAIL",
name, s.file_size_kb, "—", "—", "—"
);
}
}
}
return Ok(());
}
if backend == "topk" {
match parts.next().and_then(|s| s.parse::<usize>().ok()) {
Some(n) if n > 0 => {
self.agent.set_top_k(n);
println!("Top-k set to {}.", n);
}
_ => println!("Usage: /embed topk <N> (current: {})", self.agent.top_k()),
}
return Ok(());
}
if backend == "threshold" {
match parts.next().and_then(|s| s.parse::<f64>().ok()) {
Some(f) if f >= 0.0 => {
self.agent.set_similarity_threshold(f);
println!("Similarity threshold set to {:.3}.", f);
}
_ => println!(
"Usage: /embed threshold <F> (current: {:.3})",
self.agent.similarity_threshold()
),
}
return Ok(());
}
let model = parts.next();
let spec = match EmbedderSpec::parse(backend, model) {
Ok(s) => s,
Err(e) => {
println!("Error: {}", e);
return Ok(());
}
};
match spec.build() {
Ok(new_embedder) => {
let old_backend = self.agent.embedder().backend_name();
let old_model = self.agent.embedder().model_name().to_string();
self.agent.set_embedder(new_embedder);
println!(
"Embedder swapped: {} ({}) → {} ({})",
old_backend,
old_model,
self.agent.embedder().backend_name(),
self.agent.embedder().model_name()
);
}
Err(e) => {
println!("Failed to build embedder: {}", e);
}
}
Ok(())
}
async fn cmd_hist(&mut self, args_str: &str) -> Result<()> {
let arg = args_str.trim();
if arg.is_empty() || arg == "list" {
match self.session_store.list().await {
Ok(manifests) if manifests.is_empty() => {
println!("No saved sessions.");
}
Ok(manifests) => {
println!("{} saved session(s):", manifests.len());
for m in &manifests {
println!(
" {} — {} turns — {:?}",
m.id.0, m.turn_count, m.created
);
}
}
Err(e) => println!("Error listing sessions: {}", e),
}
return Ok(());
}
let mut parts = arg.split_whitespace();
let sub = parts.next().unwrap_or("");
let id = parts.next().unwrap_or("");
match sub {
"load" if !id.is_empty() => {
let sid = SessionId(id.to_string());
match self.session_store.load(&sid).await {
Ok(Some(session)) => {
self.prompt_memory = session.turns;
println!(
"Loaded session {} ({} turns).",
id,
self.prompt_memory.len()
);
}
Ok(None) => println!("Session '{}' not found.", id),
Err(e) => println!("Error loading session: {}", e),
}
}
"delete" if !id.is_empty() => {
let sid = SessionId(id.to_string());
match self.session_store.delete(&sid).await {
Ok(()) => println!("Deleted session '{}'.", id),
Err(e) => println!("Error deleting session: {}", e),
}
}
_ => {
println!("Usage: /hist [list | load <id> | delete <id>]");
}
}
Ok(())
}
async fn cmd_memory(&mut self, args_str: &str) -> Result<()> {
let arg = args_str.trim();
if arg.is_empty() {
let mem = if self.agent.rewriter().is_some() { "rewrite" } else { "off" };
let diff = match &self.history_strategy {
Some(s) => s.name(),
None => "off",
};
println!(
"Memory: {} — {} turns | history diffusion: {}",
mem,
self.prompt_memory.len(),
diff,
);
println!(
"Usage: /memory <backend> [model] [api_key] | transcript | log | summary | off | purge"
);
println!(" backends: ollama, deepseek");
println!(" modes: transcript — raw memory, no query rewriting");
println!(" log — enable history diffusion (raw last session)");
println!(" summary — enable history diffusion (LLM summarisation)");
return Ok(());
}
if arg.eq_ignore_ascii_case("purge") {
let count = self.prompt_memory.len();
self.prompt_memory.clear();
if let Some(ref rewriter) = self.agent.rewriter()
&& let Err(e) = rewriter.clear_memory().await {
eprintln!("Warning: memory clear failed: {}", e);
}
println!("Conversation memory purged ({} entries removed).", count);
return Ok(());
}
if arg.eq_ignore_ascii_case("off") || arg.eq_ignore_ascii_case("none") {
let was = self.agent.rewriter().is_some();
self.agent.set_rewriter(None);
if was {
println!("Memory disabled (was: rewrite)");
} else {
println!("Memory already off.");
}
return Ok(());
}
if arg.eq_ignore_ascii_case("log") {
let old = self.history_strategy.replace(Box::new(LogHistory));
match old {
Some(o) if o.name() == "log" => {
println!("History diffusion unchanged: log");
}
Some(o) => {
println!("History diffusion: {} → log", o.name());
}
None => println!("History diffusion enabled: log"),
}
return Ok(());
}
if arg.eq_ignore_ascii_case("summary") {
let summary_spec = ChatAgentSpec::ollama(
self.args.memory_model.clone(),
GenerationParams::default(),
);
match summary_spec.build() {
Ok(summary_agent) => {
let strat: Box<dyn HistoryStrategy> =
Box::new(SummaryHistory::new(summary_agent));
let old = self.history_strategy.replace(strat);
match old {
Some(o) if o.name() == "summary" => {
println!("History diffusion unchanged: summary");
}
Some(o) => {
println!("History diffusion: {} → summary", o.name());
}
None => println!("History diffusion enabled: summary"),
}
}
Err(e) => println!("Failed to build summary agent: {}", e),
}
return Ok(());
}
if arg.eq_ignore_ascii_case("transcript") {
let was = self.agent.rewriter().is_some();
self.agent.set_rewriter(None);
if was {
println!("Memory strategy: rewrite → transcript");
} else {
println!("Memory unchanged: transcript");
}
return Ok(());
}
let mut parts = arg.split_whitespace();
let backend = parts.next().unwrap_or("");
let model = parts.next();
let api_key = parts.next();
let spec = match ChatAgentSpec::parse(backend, model, api_key, None) {
Ok(s) => s,
Err(e) => {
println!("Error: {}", e);
return Ok(());
}
};
match spec.build() {
Ok(new_rewriter) => {
let new_backend = new_rewriter.backend_name();
let new_model = new_rewriter.model_name().to_string();
let was = self.agent.rewriter().is_some();
self.agent.set_rewriter(Some(new_rewriter));
if was {
println!("Memory agent: {} ({})", new_backend, new_model);
} else {
println!("Memory enabled: rewrite — {} ({})", new_backend, new_model);
}
}
Err(e) => {
println!("Failed to build memory agent: {}", e);
}
}
Ok(())
}
async fn cmd_prompt(&mut self, args_str: &str) -> Result<()> {
let mut parts = args_str.split_whitespace();
let sub = parts.next().unwrap_or("");
if sub.is_empty() {
println!("Current prompts:");
println!(
" chat (docs): {:.80}",
self.agent.system_prompt().trim()
);
println!(
" chat (no docs): {:.80}",
self.agent.chat_without_docs_prompt().trim()
);
println!(" rewrite: {:.80}", self.agent.rewrite_prompt().trim());
println!("Usage: /prompt chat|rewrite <file> or /prompt reset");
return Ok(());
}
match sub {
"reset" => {
self.agent.set_system_prompt(
"You are a helpful document assistant. Answer the user's question \
explicitly using the provided Context snippets.\n\
\n\
Context:\n{context}\n".to_string()
);
self.agent.set_rewrite_prompt(
"You are a query rewriter. Given the conversation and the \
latest question, produce a single self-contained search query \
that captures all relevant context. Output ONLY the rewritten \
query, nothing else.\n\n\
Latest question: {question}".to_string()
);
println!("Prompts reset to defaults.");
}
"chat" => {
let file = parts.next();
let Some(file) = file else {
println!("Usage: /prompt chat <file>");
return Ok(());
};
match fs::read_to_string(file) {
Ok(text) => {
self.agent.set_system_prompt(text);
println!("Chat prompt loaded from {}", file);
}
Err(e) => println!("Failed to load prompt: {}", e),
}
}
"rewrite" => {
let file = parts.next();
let Some(file) = file else {
println!("Usage: /prompt rewrite <file>");
return Ok(());
};
match fs::read_to_string(file) {
Ok(text) => {
self.agent.set_rewrite_prompt(text);
println!("Rewrite prompt loaded from {}", file);
}
Err(e) => println!("Failed to load prompt: {}", e),
}
}
other => {
println!(
"Unknown sub-command: {}. Use chat, rewrite, or reset.",
other
);
}
}
Ok(())
}
async fn cmd_parser(&mut self, args_str: &str) -> Result<()> {
let mut parts = args_str.split_whitespace();
let format = parts.next().unwrap_or("");
if format.is_empty() {
println!("PDF: {:?}", self.pdf_parser);
println!("EPUB: {:?}", self.epub_parser);
println!("Usage: /parser pdf unpdf|sink|extract|internal|vision");
println!(" /parser epub epub");
return Ok(());
}
let choice = parts.next().unwrap_or("");
if choice.is_empty() {
println!("Usage: /parser {} <backend>", format);
return Ok(());
}
match format.to_lowercase().as_str() {
"pdf" => {
#[allow(deprecated)]
let new = match choice.to_lowercase().as_str() {
#[cfg(feature = "kreuzberg")]
"kreuzberg" => PdfParserBackend::Kreuzberg,
"unpdf" => PdfParserBackend::Unpdf,
"sink" => PdfParserBackend::Sink,
"extract" => PdfParserBackend::Extract,
"internal" => PdfParserBackend::Internal,
"vision" => PdfParserBackend::Vision,
other => {
println!(
"Unknown PDF parser: {}. Use unpdf, sink, extract, internal, or vision.",
other
);
return Ok(());
}
};
let old = std::mem::replace(&mut self.pdf_parser, new.clone());
println!("PDF parser: {:?} → {:?}", old, new);
self.doc_parsers =
DocumentParsers::new(filtered_parsers(&new, self.args.sloppy_pdf));
println!("Active parsers: {}", self.doc_parsers.names().join(", "));
}
"epub" => {
let new = match choice.to_lowercase().as_str() {
"epub" => EpubParserBackend::Epub,
other => {
println!("Unknown EPUB parser: {}. The only option is 'epub'.", other);
return Ok(());
}
};
let old = std::mem::replace(&mut self.epub_parser, new.clone());
println!("EPUB parser: {:?} → {:?}", old, new);
}
other => {
println!("Unknown format: {}. Use pdf or epub.", other);
}
}
Ok(())
}
async fn cmd_rag_query(&mut self, query: &str) -> Result<()> {
let history_context = if let Some(ref strat) = self.history_strategy {
match strat.build_context(&*self.session_store, query).await {
Ok(ctx) if !ctx.is_empty() => {
eprintln!("[DEBUG] History diffusion: {} chars", ctx.len());
Some(ctx)
}
_ => None,
}
} else {
None
};
let effective_query = if let Some(ref hc) = history_context {
format!("{hc}\n\nCurrent question: {query}")
} else {
query.to_string()
};
let transcript: Vec<(&str, &str)> = self.prompt_memory.iter()
.map(|t| (t.role.as_str(), t.text.as_str()))
.collect();
eprintln!(
"[DEBUG] Provider: {} | Model: {}",
self.agent.chat_agent().backend_name(),
self.agent.chat_agent().model_name()
);
print!("Assistant > ");
stdout().flush()?;
let response_text = Arc::new(Mutex::new(String::new()));
let mut retried = false;
loop {
let rt = response_text.clone();
match self
.agent
.generate_with_context_streaming(&effective_query, &transcript, &move |text: String| {
print!("{}", text);
let _ = stdout().flush();
rt.lock().unwrap().push_str(&text);
})
.await
{
Ok(()) => {
let reply = {
let guard = response_text.lock().unwrap();
guard.clone()
};
if retried {
eprintln!(
"*** Budget permanently adjusted to {} tokens. Use `/chat context` to change. ***",
self.agent.context_tokens()
);
}
if !reply.trim().is_empty() {
self.prompt_memory.push(Turn {
role: TurnRole::User,
text: query.to_string(),
perf: None,
});
self.prompt_memory.push(Turn {
role: TurnRole::Assistant,
text: reply.trim().to_string(),
perf: None,
});
let _ = self.auto_save().await;
}
break;
}
Err(e) => {
if let Some(ce) = e.downcast_ref::<RagrigError>() {
if self.context_size_forced == ContextSizeMode::Auto && !retried {
retried = true;
let max = ce.max_size();
self.agent.set_context_tokens(max);
eprintln!(
"\n*** Context overflow: model allows {} tokens, prompt needed {}. ***",
ce.max_size(),
ce.current_size()
);
eprintln!(
"*** Budget auto‑adjusted to {} tokens — retrying. ***",
self.agent.context_tokens()
);
eprintln!(
"*** Use `/chat context {}` to override, or `--context-size-forced` to disable auto‑retry. ***",
ce.max_size().saturating_sub(512)
);
response_text.lock().unwrap().clear();
continue;
}
eprintln!(
"\n[ERROR] Prompt exceeds model context window. Model allows {} tokens, needed {}. Try `/chat context {}`.",
ce.max_size(),
ce.current_size(),
ce.max_size().saturating_sub(512)
);
} else {
eprintln!("\n[ERROR] Generation failed: {}", e);
}
break;
}
}
}
println!();
Ok(())
}
}
#[tokio::main]
async fn main() -> Result<()> {
let args = Args::parse();
let mut session = bootstrap(args).await?;
loop {
let readline = session.rl.readline("Query > ");
let cmd = match readline {
Ok(line) => {
let trimmed = line.trim().to_string();
if trimmed.is_empty() {
continue;
}
session.rl.add_history_entry(&trimmed)?;
Command::from(trimmed.as_str())
}
Err(ReadlineError::Interrupted) => {
println!("\nSession interrupted.");
break;
}
Err(ReadlineError::Eof) => {
println!("\nSession ended via Ctrl+D.");
break;
}
Err(err) => {
eprintln!("Error reading input: {}", err);
break;
}
};
match cmd {
Command::Exit => break,
Command::RagQuery(ref q) if q.is_empty() => continue,
_ => {
if let Err(e) = session.execute(cmd).await {
eprintln!("Error: {}", e);
}
}
}
}
if !session.prompt_memory.is_empty()
&& let Err(e) = session.auto_save().await {
eprintln!("Warning: failed to save session: {}", e);
}
session.rl.save_history(&session.history_path)?;
Ok(())
}
fn strip_ansi(input: &str) -> String {
let mut result = String::with_capacity(input.len());
let mut chars = input.chars().peekable();
while let Some(c) = chars.next() {
if c == '\x1b' && chars.peek() == Some(&'[') {
chars.next(); while let Some(&nc) = chars.peek() {
chars.next();
if nc.is_alphabetic() || nc == '~' {
break;
}
}
} else {
result.push(c);
}
}
result
}
fn parse_number_range(input: &str) -> Result<Vec<usize>, String> {
let mut indices = Vec::new();
for part in input.split(',') {
let part = part.trim();
if part.is_empty() {
continue;
}
if let Some((start, end)) = part.split_once('-') {
let s: usize = start
.trim()
.parse()
.map_err(|_| format!("invalid number: {}", start))?;
let e: usize = end
.trim()
.parse()
.map_err(|_| format!("invalid number: {}", end))?;
if s == 0 || e == 0 {
return Err("Indices start at 1".to_string());
}
if s > e {
return Err(format!("invalid range: {}-{}", s, e));
}
for n in s..=e {
indices.push(n - 1); }
} else {
let n: usize = part
.parse()
.map_err(|_| format!("invalid number: {}", part))?;
if n == 0 {
return Err("Indices start at 1".to_string());
}
indices.push(n - 1);
}
}
Ok(indices)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_unknown_slash_command_is_not_query() {
let cmd = Command::from("/foobar");
assert!(matches!(cmd, Command::Unknown(_)));
}
#[test]
fn parse_unknown_slash_with_args() {
let cmd = Command::from("/bogus arg1 arg2");
assert!(matches!(cmd, Command::Unknown(c) if c.contains("arg1")));
}
#[test]
fn parse_plain_text_is_rag_query() {
let cmd = Command::from("What is RAG?");
assert!(matches!(cmd, Command::RagQuery(q) if q == "What is RAG?"));
}
#[test]
fn parse_memory_no_args_does_not_panic() {
let cmd = Command::from("/memory");
assert!(matches!(cmd, Command::Memory(s) if s.is_empty()));
}
#[test]
fn parse_memory_with_args() {
let cmd = Command::from("/memory transcript");
assert!(matches!(cmd, Command::Memory(s) if s == "transcript"));
}
#[test]
fn parse_hist_no_args_does_not_panic() {
let cmd = Command::from("/hist");
assert!(matches!(cmd, Command::Hist(s) if s.is_empty()));
}
#[test]
fn parse_chat_command_recognised() {
let cmd = Command::from("/chat ollama");
assert!(matches!(cmd, Command::Chat(s) if s == "ollama"));
}
#[test]
fn parse_embed_no_args_does_not_panic() {
let cmd = Command::from("/embed");
assert!(matches!(cmd, Command::Embed(s) if s.is_empty()));
}
#[test]
fn parse_refs_no_args_does_not_panic() {
let cmd = Command::from("/refs");
assert!(matches!(cmd, Command::ExtractRefs(s) if s.is_empty()));
}
#[test]
fn parse_slash_exit() {
assert!(matches!(Command::from("/exit"), Command::Exit));
}
#[test]
fn parse_slash_bye() {
assert!(matches!(Command::from("/bye"), Command::Exit));
}
#[tokio::test]
#[ignore = "requires Ollama with gemma4:e4b and tests/fixtures/formats/pdf"]
async fn gemma4_rag_answer_exceeds_20_words() {
let args = Args::parse_from([
"test",
"--folder",
"tests/fixtures/formats/pdf",
"--model",
"gemma4:e4b",
"--embedding-model",
"nomic-embed-text",
]);
let mut session = match bootstrap(args).await {
Ok(s) => s,
Err(e) => {
eprintln!("bootstrap failed (Ollama not running?): {}", e);
return;
}
};
let question = "I have used a 7-item Likert scale in my research. What should I do?";
match session.cmd_rag_query(question).await {
Ok(()) => {
let memory = &session.prompt_memory;
let answer = memory
.last()
.filter(|t| t.role == TurnRole::Assistant)
.map(|t| t.text.as_str())
.unwrap_or("");
let word_count = answer.split_whitespace().count();
eprintln!("Answer ({} words): {}", word_count, answer);
assert!(
word_count > 20,
"Expected >20 words, got {}: '{}'",
word_count,
answer
);
}
Err(e) => {
eprintln!("RAG query failed: {}", e);
panic!("cmd_rag_query returned error: {}", e);
}
}
}
}