use anyhow::{Context, Result};
use clap::Parser;
use frozen_duckdb::cli::commands::{Cli, Commands};
use frozen_duckdb::cli::dataset_manager::DatasetManager;
use frozen_duckdb::cli::flock_manager::FlockManager;
use serde_json::{self, Value};
use std::io;
use std::path::Path;
use tracing::{error, info};
fn main() -> Result<()> {
let cli = Cli::parse();
let subscriber = tracing_subscriber::FmtSubscriber::builder()
.with_max_level(match cli.verbose {
0 => tracing::Level::WARN,
1 => tracing::Level::INFO,
2 => tracing::Level::DEBUG,
_ => tracing::Level::TRACE,
})
.finish();
tracing::subscriber::set_global_default(subscriber).expect("Failed to set tracing subscriber");
match cli.command {
Commands::Download {
dataset,
output_dir,
format,
} => {
let dataset_manager = DatasetManager::new()?;
match dataset.as_str() {
"chinook" => {
dataset_manager.download_chinook(&output_dir, &format)?;
}
"tpch" => {
dataset_manager.download_tpch(&output_dir, &format)?;
}
_ => {
error!("โ Unknown dataset: {}", dataset);
error!(" Available datasets: chinook, tpch");
std::process::exit(1);
}
}
}
Commands::Convert {
input,
output,
input_format,
output_format,
} => {
let dataset_manager = DatasetManager::new()?;
dataset_manager.convert_dataset(&input, &output, &input_format, &output_format)?;
}
Commands::Info => {
let dataset_manager = DatasetManager::new()?;
dataset_manager.show_info()?;
}
Commands::FlockSetup {
ollama_url,
text_model,
embedding_model,
skip_verification,
} => {
let flock_manager = FlockManager::new()?;
if !flock_manager.is_flock_ready()? {
error!("โ Flock extension not available");
error!(" Make sure DuckDB with Flock extension is properly installed");
std::process::exit(4);
}
flock_manager.setup_ollama(&ollama_url, &text_model, &embedding_model, skip_verification)?;
}
Commands::Complete {
prompt,
input,
output,
model,
max_tokens: _,
temperature: _,
} => {
let flock_manager = FlockManager::new()?;
if !flock_manager.is_flock_ready()? {
error!("โ Flock extension not available");
error!(" Run 'frozen-duckdb flock-setup' first");
std::process::exit(4);
}
let text_to_complete = if let Some(prompt_text) = prompt {
prompt_text
} else if let Some(input_file) = input {
match std::fs::read_to_string(&input_file) {
Ok(content) => content.trim().to_string(),
Err(e) => {
error!("โ Failed to read input file '{}': {}", input_file, e);
std::process::exit(1);
}
}
} else {
info!("๐ Enter text to complete (Ctrl+D to finish):");
let mut buffer = String::new();
io::stdin().read_line(&mut buffer)?;
buffer.trim().to_string()
};
let response = flock_manager.complete_text(&text_to_complete, model.as_str())
.unwrap_or_else(|_| {
error!("โ Text completion failed - check if Ollama is running");
std::process::exit(1);
});
if let Some(output_file) = output {
match std::fs::write(&output_file, &response) {
Ok(_) => info!("โ
Response written to: {}", output_file),
Err(e) => {
error!("โ Failed to write to output file '{}': {}", output_file, e);
std::process::exit(1);
}
}
} else {
println!("{}", response);
}
}
Commands::Embed {
text,
input,
output,
model,
normalize,
} => {
let flock_manager = FlockManager::new()?;
if !flock_manager.is_flock_ready()? {
error!("โ Flock extension not available");
error!(" Run 'frozen-duckdb flock-setup' first");
std::process::exit(4);
}
let texts_to_embed = if let Some(text_content) = text {
vec![text_content.clone()]
} else if let Some(input_file) = input {
match std::fs::read_to_string(&input_file) {
Ok(content) => content.lines().map(|s| s.to_string()).collect(),
Err(e) => {
error!("โ Failed to read input file '{}': {}", input_file, e);
std::process::exit(1);
}
}
} else {
error!("โ Must provide either --text or --input");
std::process::exit(1);
};
let embeddings = flock_manager.generate_embeddings(texts_to_embed, &model, normalize)
.expect("Embedding generation not implemented yet");
if let Some(output_file) = output {
let json_data = serde_json::to_string_pretty(&embeddings)
.context("Failed to serialize embeddings to JSON")?;
match std::fs::write(&output_file, json_data) {
Ok(_) => info!("โ
Embeddings written to: {}", output_file),
Err(e) => {
error!("โ Failed to write to output file '{}': {}", output_file, e);
std::process::exit(1);
}
}
} else {
println!("{}", serde_json::to_string_pretty(&embeddings)?);
}
}
Commands::Search {
query,
corpus,
threshold,
limit,
format,
} => {
let flock_manager = FlockManager::new()?;
if !flock_manager.is_flock_ready()? {
error!("โ Flock extension not available");
error!(" Run 'frozen-duckdb flock-setup' first");
std::process::exit(4);
}
let results = flock_manager.semantic_search(&query, &corpus, threshold, limit)
.expect("Semantic search not implemented yet");
match format.as_str() {
"json" => {
let json_results: Vec<Value> = results
.into_iter()
.map(|(doc, score)| {
serde_json::json!({
"document": doc,
"similarity_score": score
})
})
.collect();
println!("{}", serde_json::to_string_pretty(&json_results)?);
}
_ => {
if results.is_empty() {
info!("๐ No similar documents found above threshold {:.3}", threshold);
} else {
info!("๐ Found {} similar documents:", results.len());
for (i, (doc, score)) in results.iter().enumerate() {
println!(" {}. \"{}\" (similarity: {:.3})", i + 1, doc, score);
}
}
}
}
}
Commands::Filter {
criteria,
prompt,
input,
output,
model,
positive_only,
} => {
let flock_manager = FlockManager::new()?;
if !flock_manager.is_flock_ready()? {
error!("โ Flock extension not available");
error!(" Run 'frozen-duckdb flock-setup' first");
std::process::exit(4);
}
let filter_criteria = if let Some(custom_prompt) = prompt {
custom_prompt.clone()
} else if let Some(criteria_text) = criteria {
format!("{} Answer yes or no: {{text}}", criteria_text)
} else {
error!("โ Must provide either --criteria or --prompt");
std::process::exit(1);
};
let results = flock_manager.llm_filter(&filter_criteria, &input, &model, true)
.expect("LLM filtering not implemented yet");
if let Some(output_file) = output {
let json_data = serde_json::to_string_pretty(&results)
.context("Failed to serialize filter results to JSON")?;
match std::fs::write(&output_file, json_data) {
Ok(_) => info!("โ
Filter results written to: {}", output_file),
Err(e) => {
error!("โ Failed to write to output file '{}': {}", output_file, e);
std::process::exit(1);
}
}
} else {
if positive_only {
info!("โ
Items that match criteria:");
for (item, matches) in results {
if matches {
println!("โ
{}", item);
}
}
} else {
info!("๐ Filter results:");
for (item, matches) in results {
let status = if matches { "โ
MATCH" } else { "โ NO MATCH" };
println!("{}: {}", status, item);
}
}
}
}
Commands::Summarize {
input,
output,
strategy,
max_length,
model,
} => {
let flock_manager = FlockManager::new()?;
if !flock_manager.is_flock_ready()? {
error!("โ Flock extension not available");
error!(" Run 'frozen-duckdb flock-setup' first");
std::process::exit(4);
}
let texts = if Path::new(&input).is_dir() {
let mut all_texts = Vec::new();
for entry in std::fs::read_dir(&input)? {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("txt") {
if let Ok(content) = std::fs::read_to_string(&path) {
all_texts.push(content.trim().to_string());
}
}
}
all_texts
} else {
match std::fs::read_to_string(&input) {
Ok(content) => content.lines().map(|s| s.to_string()).collect(),
Err(e) => {
error!("โ Failed to read input file '{}': {}", input, e);
std::process::exit(1);
}
}
};
let summary = flock_manager.summarize_texts(texts, &strategy, max_length, &model)
.expect("Text summarization not implemented yet");
if let Some(output_file) = output {
match std::fs::write(&output_file, &summary) {
Ok(_) => info!("โ
Summary written to: {}", output_file),
Err(e) => {
error!("โ Failed to write to output file '{}': {}", output_file, e);
std::process::exit(1);
}
}
} else {
println!("{}", summary);
}
}
Commands::Test => {
info!("๐งช Tests have been moved to the test suite");
info!(" Run tests with: cargo test");
info!(" Run specific tests with: cargo test <test_name>");
info!(" Run all tests with: cargo test --all");
}
Commands::Benchmark {
operation,
iterations,
size,
} => {
info!(
"Benchmarking {} operation with {} iterations (size: {})",
operation, iterations, size
);
info!("๐ Performance benchmarking feature coming soon!");
}
}
Ok(())
}