use anyhow::Result;
use clap::Parser;
use rig_core::client::{CompletionClient, EmbeddingsClient, Nothing};
use rig_core::completion::Prompt;
use rig_core::embeddings::EmbeddingsBuilder;
use rig_core::providers::{deepseek, ollama};
use rustyline::error::ReadlineError;
use rustyline::DefaultEditor;
use std::io::{stdout, Write};
use std::path::PathBuf;
#[derive(Parser, Debug)]
#[command(about = "Pure-rig-core RAG demo — shows what rig-core gives natively vs. what ragrig adds")]
pub struct Args {
#[arg(short, long)]
pub folder: PathBuf,
#[arg(short, long, default_value = "gemma2:latest")]
pub model: String,
#[arg(short = 'e', long, default_value = "nomic-embed-text")]
pub embedding_model: String,
#[arg(long, env = "DEEPSEEK_API_KEY")]
pub deepseek_api_key: Option<String>,
#[arg(long, default_value = "deepseek-chat")]
pub deepseek_model: String,
#[arg(long)]
pub temperature: Option<f64>,
#[arg(long)]
pub max_tokens: Option<usize>,
#[arg(long, default_value = "4096")]
pub model_ctx_tokens: usize,
}
struct Session {
args: Args,
chat_model: String,
chat_backend: String,
deepseek_api_key: Option<String>,
embed_model: String,
rl: DefaultEditor,
history_path: PathBuf,
transcript: Vec<(String, String)>,
}
enum Command {
Chat(String),
Embed(String),
Help,
Memory(String),
RagQuery(String),
Unknown(String),
Exit,
}
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("/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());
}
Command::Unknown(input.to_string())
}
}
async fn bootstrap(args: Args) -> Result<Session> {
let chat_backend = if args.deepseek_api_key.is_some() {
"deepseek".to_string()
} else {
"ollama".to_string()
};
let chat_model = if args.deepseek_api_key.is_some() {
args.deepseek_model.clone()
} else {
args.model.clone()
};
let embed_model = args.embedding_model.clone();
let deepseek_api_key = args.deepseek_api_key.clone();
println!("Chat: {} ({})", chat_backend, chat_model);
println!("Embed: ollama ({})", embed_model);
println!(
"NOTE: No documents loaded — rig-core has no vector store, document parser, or chunker."
);
println!(" ragrig provides all of these (see real main.rs).");
let mut rl = DefaultEditor::new()?;
let history_path = args.folder.join(".ragrig_history");
if history_path.exists() {
let _ = rl.load_history(&history_path);
}
println!("\nRAG System Online (rig-core only — most features stubbed).");
println!("Commands: /chat ollama|deepseek [model] [key] | /embed <model> | /help | exit");
println!("Ask questions (but there's no document index, so results come from model knowledge only):");
Ok(Session {
args,
chat_model,
chat_backend,
deepseek_api_key,
embed_model,
rl,
history_path,
transcript: Vec::new(),
})
}
impl Session {
async fn execute(&mut self, cmd: Command) -> Result<()> {
match cmd {
Command::Chat(args_str) => self.cmd_chat(&args_str).await,
Command::Embed(args_str) => self.cmd_embed(&args_str).await,
Command::Help => {
self.cmd_help();
Ok(())
}
Command::Memory(args_str) => self.cmd_memory(&args_str).await,
Command::RagQuery(q) => self.cmd_rag_query(&q).await,
Command::Unknown(cmd) => {
println!("Unknown command: '{}'", cmd);
Ok(())
}
Command::Exit => Ok(()),
}
}
fn cmd_help(&self) {
println!("/chat ollama|deepseek [model] [api_key] — hot-swap chat backend");
println!(" /chat temperature <F> — set sampling temperature");
println!(" /chat max_tokens <N> — set max output tokens");
println!("/embed <model> — change embedding model");
println!("/memory — show conversation transcript (raw, no rewriting)");
println!("/help — this message");
println!("exit / quit — end session");
println!();
println!("Any other text is treated as a query and sent to the chat model.");
println!("(There is no document index — rig-core has no vector store.)");
}
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 == "temperature" {
match parts.next().and_then(|s| s.parse::<f64>().ok()) {
Some(t) if t >= 0.0 => {
println!("Temperature set to {:.2}.", t);
eprintln!(
"[NOTE] rig-core applies temperature via .temperature() on the \
agent builder, but there is no stored agent object to mutate. \
ragrig's RagAgent::set_chat_agent() rebuilds with new params."
);
}
_ => println!("Usage: /chat temperature <F> (0.0 = deterministic)"),
}
return Ok(());
}
if backend == "max_tokens" {
match parts.next().and_then(|s| s.parse::<usize>().ok()) {
Some(n) if n > 0 => {
println!("Max tokens set to {}.", n);
}
_ => println!("Usage: /chat max_tokens <N>"),
}
return Ok(());
}
if backend.is_empty() {
println!(
"Chat: {} ({})",
self.chat_backend, self.chat_model
);
println!("Usage: /chat ollama|deepseek [model] [api_key] | temperature <F> | max_tokens <N>");
return Ok(());
}
let model = parts.next().unwrap_or("default").to_string();
let api_key = parts.next().map(|s| s.to_string());
match backend.to_lowercase().as_str() {
"ollama" => {
let old_backend = self.chat_backend.clone();
let old_model = self.chat_model.clone();
self.chat_backend = "ollama".to_string();
self.chat_model = model;
println!(
"Chat agent swapped: {} ({}) → ollama ({})",
old_backend, old_model, self.chat_model
);
}
"deepseek" => {
let key = api_key.or_else(|| self.deepseek_api_key.clone());
let Some(ref key) = key else {
println!("DeepSeek API key required. Set DEEPSEEK_API_KEY env var or pass as argument.");
return Ok(());
};
let old_backend = self.chat_backend.clone();
let old_model = self.chat_model.clone();
self.chat_backend = "deepseek".to_string();
self.chat_model = model;
self.deepseek_api_key = Some(key.clone());
println!(
"Chat agent swapped: {} ({}) → deepseek ({})",
old_backend, old_model, self.chat_model
);
}
other => {
println!(
"Unknown backend: {}. Use ollama or deepseek.",
other
);
}
}
Ok(())
}
async fn cmd_embed(&mut self, args_str: &str) -> Result<()> {
let arg = args_str.trim();
if arg.is_empty() {
println!(
"Embed: ollama ({})",
self.embed_model
);
println!("Usage: /embed <model>");
return Ok(());
}
let old = self.embed_model.clone();
self.embed_model = arg.to_string();
println!("Embedder swapped: ollama ({}) → ollama ({})", old, self.embed_model);
Ok(())
}
async fn cmd_memory(&mut self, args_str: &str) -> Result<()> {
let arg = args_str.trim();
if arg.is_empty() {
println!("Transcript: {} turns", self.transcript.len());
for (i, (role, text)) in self.transcript.iter().enumerate() {
println!(
" [{:2}] {}: {:.80}",
i + 1,
role,
text.lines().next().unwrap_or("")
);
}
println!();
println!("(ragrig provides query rewriting, history diffusion, and session persistence)");
return Ok(());
}
if arg.eq_ignore_ascii_case("purge") {
let count = self.transcript.len();
self.transcript.clear();
println!("Conversation memory cleared ({} turns removed).", count);
return Ok(());
}
println!("Unknown memory mode: {}. Use /memory or /memory purge.", arg);
Ok(())
}
async fn cmd_rag_query(&mut self, query: &str) -> Result<()> {
let query_embedding = match embed_text(&self.embed_model, &[query]).await {
Ok(embs) => embs,
Err(e) => {
eprintln!("[ERROR] Embedding failed: {}", e);
return Ok(());
}
};
eprintln!(
"[INFO] Query embedded: {} dimensions",
query_embedding.first().map(|v| v.len()).unwrap_or(0)
);
let search_results: Vec<&str> = vec![];
if search_results.is_empty() {
eprintln!(
"[NOTE] No document results — rig-core has no vector store. \
ragrig provides BruteForceStore (BM25 + cosine + RRF fusion)."
);
}
let mut prompt = String::new();
prompt.push_str(
"You are a helpful assistant. Answer the user's question based on your knowledge.\n\n",
);
for (role, text) in &self.transcript {
prompt.push_str(&format!(
"{}: {}\n",
role,
text
));
}
prompt.push_str(&format!("User: {query}\nAssistant: "));
eprintln!(
"[INFO] Generating with {} ({}) — prompt is {} chars",
self.chat_backend,
self.chat_model,
prompt.len()
);
print!("Assistant > ");
stdout().flush()?;
let reply = match self.chat_backend.as_str() {
"ollama" => {
let client = ollama::Client::new(Nothing)
.map_err(|e| anyhow::anyhow!("Ollama: {}", e))?;
let mut builder = client.agent(&self.chat_model);
if let Some(t) = self.args.temperature {
builder = builder.temperature(t);
}
let agent = builder.build();
match agent.prompt(&prompt).await {
Ok(r) => r,
Err(e) => {
eprintln!("\n[ERROR] Generation failed: {}", e);
return Ok(());
}
}
}
"deepseek" => {
let Some(ref key) = self.deepseek_api_key else {
eprintln!("\n[ERROR] DeepSeek API key not set.");
return Ok(());
};
let client = deepseek::Client::new(key)
.map_err(|e| anyhow::anyhow!("DeepSeek: {}", e))?;
let agent = client.agent(&self.chat_model).build();
match agent.prompt(&prompt).await {
Ok(r) => r,
Err(e) => {
eprintln!("\n[ERROR] Generation failed: {}", e);
return Ok(());
}
}
}
other => {
eprintln!("\n[ERROR] Unknown backend: {}", other);
return Ok(());
}
};
print!("{}", reply);
println!();
if !reply.trim().is_empty() {
self.transcript.push(("User".to_string(), query.to_string()));
self.transcript.push(("Assistant".to_string(), reply.trim().to_string()));
}
Ok(())
}
}
async fn embed_text(model: &str, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let client = ollama::Client::new(Nothing)
.map_err(|e| anyhow::anyhow!("Ollama embedder: {}", e))?;
let embed_model = client.embedding_model(model);
let owned: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
let embedded = EmbeddingsBuilder::new(embed_model)
.documents(owned)?
.build()
.await
.map_err(|e| anyhow::anyhow!("Embedding failed: {}", e))?;
Ok(embedded
.into_iter()
.map(|(_text, emb)| emb.first().vec.iter().map(|v| *v as f32).collect())
.collect())
}
#[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);
}
}
}
}
let _ = session.rl.save_history(&session.history_path);
Ok(())
}