Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'
This commit is contained in:
344
vendor/ruvector/crates/ruvector-cli/src/cli/commands.rs
vendored
Normal file
344
vendor/ruvector/crates/ruvector-cli/src/cli/commands.rs
vendored
Normal file
@@ -0,0 +1,344 @@
|
||||
//! CLI command implementations
|
||||
|
||||
use crate::cli::{
|
||||
export_csv, export_json, format_error, format_search_results, format_stats, format_success,
|
||||
ProgressTracker,
|
||||
};
|
||||
use crate::config::Config;
|
||||
use anyhow::{Context, Result};
|
||||
use colored::*;
|
||||
use ruvector_core::{
|
||||
types::{DbOptions, SearchQuery, VectorEntry},
|
||||
VectorDB,
|
||||
};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Instant;
|
||||
|
||||
/// Create a new database
|
||||
pub fn create_database(path: &str, dimensions: usize, config: &Config) -> Result<()> {
|
||||
let mut db_options = config.to_db_options();
|
||||
db_options.storage_path = path.to_string();
|
||||
db_options.dimensions = dimensions;
|
||||
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Creating database at: {}", path))
|
||||
);
|
||||
println!(" Dimensions: {}", dimensions.to_string().cyan());
|
||||
println!(" Distance metric: {:?}", db_options.distance_metric);
|
||||
|
||||
let _db = VectorDB::new(db_options).context("Failed to create database")?;
|
||||
|
||||
println!("{}", format_success("Database created successfully!"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Insert vectors from a file
|
||||
pub fn insert_vectors(
|
||||
db_path: &str,
|
||||
input_file: &str,
|
||||
format: &str,
|
||||
config: &Config,
|
||||
show_progress: bool,
|
||||
) -> Result<()> {
|
||||
// Load database
|
||||
let mut db_options = config.to_db_options();
|
||||
db_options.storage_path = db_path.to_string();
|
||||
|
||||
let db = VectorDB::new(db_options).context("Failed to open database")?;
|
||||
|
||||
// Parse input file
|
||||
let entries = match format {
|
||||
"json" => parse_json_file(input_file)?,
|
||||
"csv" => parse_csv_file(input_file)?,
|
||||
"npy" => parse_npy_file(input_file)?,
|
||||
_ => return Err(anyhow::anyhow!("Unsupported format: {}", format)),
|
||||
};
|
||||
|
||||
let total = entries.len();
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Loaded {} vectors from {}", total, input_file))
|
||||
);
|
||||
|
||||
// Insert with progress
|
||||
let start = Instant::now();
|
||||
let tracker = ProgressTracker::new();
|
||||
let pb = if show_progress {
|
||||
Some(tracker.create_bar(total as u64, "Inserting vectors..."))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let batch_size = config.cli.batch_size;
|
||||
let mut inserted = 0;
|
||||
|
||||
for chunk in entries.chunks(batch_size) {
|
||||
db.insert_batch(chunk.to_vec())
|
||||
.context("Failed to insert batch")?;
|
||||
inserted += chunk.len();
|
||||
|
||||
if let Some(ref pb) = pb {
|
||||
pb.set_position(inserted as u64);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pb) = pb {
|
||||
pb.finish_with_message("Insertion complete!");
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!(
|
||||
"Inserted {} vectors in {:.2}s ({:.0} vectors/sec)",
|
||||
total,
|
||||
elapsed.as_secs_f64(),
|
||||
total as f64 / elapsed.as_secs_f64()
|
||||
))
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Search for similar vectors
|
||||
pub fn search_vectors(
|
||||
db_path: &str,
|
||||
query_vector: Vec<f32>,
|
||||
k: usize,
|
||||
config: &Config,
|
||||
show_vectors: bool,
|
||||
) -> Result<()> {
|
||||
let mut db_options = config.to_db_options();
|
||||
db_options.storage_path = db_path.to_string();
|
||||
|
||||
let db = VectorDB::new(db_options).context("Failed to open database")?;
|
||||
|
||||
let start = Instant::now();
|
||||
let results = db
|
||||
.search(SearchQuery {
|
||||
vector: query_vector,
|
||||
k,
|
||||
filter: None,
|
||||
ef_search: None,
|
||||
})
|
||||
.context("Failed to search")?;
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
println!("{}", format_search_results(&results, show_vectors));
|
||||
println!(
|
||||
"\n{}",
|
||||
format!(
|
||||
"Search completed in {:.2}ms",
|
||||
elapsed.as_secs_f64() * 1000.0
|
||||
)
|
||||
.dimmed()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Show database information
|
||||
pub fn show_info(db_path: &str, config: &Config) -> Result<()> {
|
||||
let mut db_options = config.to_db_options();
|
||||
db_options.storage_path = db_path.to_string();
|
||||
|
||||
let db = VectorDB::new(db_options).context("Failed to open database")?;
|
||||
|
||||
let count = db.len().context("Failed to get count")?;
|
||||
let dimensions = db.options().dimensions;
|
||||
let metric = format!("{:?}", db.options().distance_metric);
|
||||
|
||||
println!("{}", format_stats(count, dimensions, &metric));
|
||||
|
||||
if let Some(hnsw_config) = &db.options().hnsw_config {
|
||||
println!("{}", "HNSW Configuration:".bold().green());
|
||||
println!(" M: {}", hnsw_config.m.to_string().cyan());
|
||||
println!(
|
||||
" ef_construction: {}",
|
||||
hnsw_config.ef_construction.to_string().cyan()
|
||||
);
|
||||
println!(" ef_search: {}", hnsw_config.ef_search.to_string().cyan());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run a quick benchmark
|
||||
pub fn run_benchmark(db_path: &str, config: &Config, num_queries: usize) -> Result<()> {
|
||||
let mut db_options = config.to_db_options();
|
||||
db_options.storage_path = db_path.to_string();
|
||||
|
||||
let db = VectorDB::new(db_options).context("Failed to open database")?;
|
||||
|
||||
let dimensions = db.options().dimensions;
|
||||
|
||||
println!("{}", "Running benchmark...".bold().green());
|
||||
println!(" Queries: {}", num_queries.to_string().cyan());
|
||||
println!(" Dimensions: {}", dimensions.to_string().cyan());
|
||||
|
||||
// Generate random query vectors
|
||||
use rand::Rng;
|
||||
let mut rng = rand::thread_rng();
|
||||
let queries: Vec<Vec<f32>> = (0..num_queries)
|
||||
.map(|_| (0..dimensions).map(|_| rng.gen()).collect())
|
||||
.collect();
|
||||
|
||||
// Warm-up
|
||||
for query in queries.iter().take(10) {
|
||||
let _ = db.search(SearchQuery {
|
||||
vector: query.clone(),
|
||||
k: 10,
|
||||
filter: None,
|
||||
ef_search: None,
|
||||
});
|
||||
}
|
||||
|
||||
// Benchmark
|
||||
let start = Instant::now();
|
||||
for query in &queries {
|
||||
db.search(SearchQuery {
|
||||
vector: query.clone(),
|
||||
k: 10,
|
||||
filter: None,
|
||||
ef_search: None,
|
||||
})
|
||||
.context("Search failed")?;
|
||||
}
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
let qps = num_queries as f64 / elapsed.as_secs_f64();
|
||||
let avg_latency = elapsed.as_secs_f64() * 1000.0 / num_queries as f64;
|
||||
|
||||
println!("\n{}", "Benchmark Results:".bold().green());
|
||||
println!(" Total time: {:.2}s", elapsed.as_secs_f64());
|
||||
println!(" Queries per second: {:.0}", qps.to_string().cyan());
|
||||
println!(" Average latency: {:.2}ms", avg_latency.to_string().cyan());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Export database to file
|
||||
pub fn export_database(
|
||||
db_path: &str,
|
||||
output_file: &str,
|
||||
format: &str,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
let mut db_options = config.to_db_options();
|
||||
db_options.storage_path = db_path.to_string();
|
||||
|
||||
let db = VectorDB::new(db_options).context("Failed to open database")?;
|
||||
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Exporting database to: {}", output_file))
|
||||
);
|
||||
|
||||
// Export is currently limited - would need to add all_ids() method to VectorDB
|
||||
// For now, return an error with a helpful message
|
||||
return Err(anyhow::anyhow!(
|
||||
"Export functionality requires VectorDB::all_ids() method. This will be implemented in a future update."
|
||||
));
|
||||
|
||||
// TODO: Implement when VectorDB exposes all_ids()
|
||||
// let ids = db.all_ids()?;
|
||||
// let tracker = ProgressTracker::new();
|
||||
// let pb = tracker.create_bar(ids.len() as u64, "Exporting vectors...");
|
||||
// ...
|
||||
}
|
||||
|
||||
/// Import from other vector databases
|
||||
pub fn import_from_external(
|
||||
db_path: &str,
|
||||
source: &str,
|
||||
source_path: &str,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Importing from {} database", source))
|
||||
);
|
||||
|
||||
match source {
|
||||
"faiss" => {
|
||||
// TODO: Implement FAISS import
|
||||
return Err(anyhow::anyhow!("FAISS import not yet implemented"));
|
||||
}
|
||||
"pinecone" => {
|
||||
// TODO: Implement Pinecone import
|
||||
return Err(anyhow::anyhow!("Pinecone import not yet implemented"));
|
||||
}
|
||||
"weaviate" => {
|
||||
// TODO: Implement Weaviate import
|
||||
return Err(anyhow::anyhow!("Weaviate import not yet implemented"));
|
||||
}
|
||||
_ => return Err(anyhow::anyhow!("Unsupported source: {}", source)),
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
fn parse_json_file(path: &str) -> Result<Vec<VectorEntry>> {
|
||||
let content = std::fs::read_to_string(path).context("Failed to read JSON file")?;
|
||||
serde_json::from_str(&content).context("Failed to parse JSON")
|
||||
}
|
||||
|
||||
fn parse_csv_file(path: &str) -> Result<Vec<VectorEntry>> {
|
||||
let mut reader = csv::Reader::from_path(path).context("Failed to open CSV file")?;
|
||||
|
||||
let mut entries = Vec::new();
|
||||
|
||||
for result in reader.records() {
|
||||
let record = result.context("Failed to read CSV record")?;
|
||||
|
||||
let id = if record.get(0).map(|s| s.is_empty()).unwrap_or(true) {
|
||||
None
|
||||
} else {
|
||||
Some(record.get(0).unwrap().to_string())
|
||||
};
|
||||
|
||||
let vector: Vec<f32> =
|
||||
serde_json::from_str(record.get(1).context("Missing vector column")?)
|
||||
.context("Failed to parse vector")?;
|
||||
|
||||
let metadata = if let Some(meta_str) = record.get(2) {
|
||||
if !meta_str.is_empty() {
|
||||
Some(serde_json::from_str(meta_str).context("Failed to parse metadata")?)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
entries.push(VectorEntry {
|
||||
id,
|
||||
vector,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
fn parse_npy_file(path: &str) -> Result<Vec<VectorEntry>> {
|
||||
use ndarray::Array2;
|
||||
use ndarray_npy::ReadNpyExt;
|
||||
|
||||
let file = std::fs::File::open(path).context("Failed to open NPY file")?;
|
||||
let array: Array2<f32> = Array2::read_npy(file).context("Failed to read NPY file")?;
|
||||
|
||||
let entries: Vec<VectorEntry> = array
|
||||
.outer_iter()
|
||||
.enumerate()
|
||||
.map(|(i, row)| VectorEntry {
|
||||
id: Some(format!("vec_{}", i)),
|
||||
vector: row.to_vec(),
|
||||
metadata: None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(entries)
|
||||
}
|
||||
179
vendor/ruvector/crates/ruvector-cli/src/cli/format.rs
vendored
Normal file
179
vendor/ruvector/crates/ruvector-cli/src/cli/format.rs
vendored
Normal file
@@ -0,0 +1,179 @@
|
||||
//! Output formatting utilities
|
||||
|
||||
use colored::*;
|
||||
use ruvector_core::types::{SearchResult, VectorEntry};
|
||||
use serde_json;
|
||||
|
||||
/// Format search results for display
|
||||
pub fn format_search_results(results: &[SearchResult], show_vectors: bool) -> String {
|
||||
let mut output = String::new();
|
||||
|
||||
for (i, result) in results.iter().enumerate() {
|
||||
output.push_str(&format!("\n{}. {}\n", i + 1, result.id.bold()));
|
||||
output.push_str(&format!(" Score: {:.4}\n", result.score));
|
||||
|
||||
if let Some(metadata) = &result.metadata {
|
||||
if !metadata.is_empty() {
|
||||
output.push_str(&format!(
|
||||
" Metadata: {}\n",
|
||||
serde_json::to_string_pretty(metadata).unwrap_or_else(|_| "{}".to_string())
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if show_vectors {
|
||||
if let Some(vector) = &result.vector {
|
||||
let preview: Vec<f32> = vector.iter().take(5).copied().collect();
|
||||
output.push_str(&format!(" Vector (first 5): {:?}...\n", preview));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// Format database statistics
|
||||
pub fn format_stats(count: usize, dimensions: usize, metric: &str) -> String {
|
||||
format!(
|
||||
"\n{}\n Vectors: {}\n Dimensions: {}\n Distance Metric: {}\n",
|
||||
"Database Statistics".bold().green(),
|
||||
count.to_string().cyan(),
|
||||
dimensions.to_string().cyan(),
|
||||
metric.cyan()
|
||||
)
|
||||
}
|
||||
|
||||
/// Format error message
|
||||
pub fn format_error(msg: &str) -> String {
|
||||
format!("{} {}", "Error:".red().bold(), msg)
|
||||
}
|
||||
|
||||
/// Format success message
|
||||
pub fn format_success(msg: &str) -> String {
|
||||
format!("{} {}", "✓".green().bold(), msg)
|
||||
}
|
||||
|
||||
/// Format warning message
|
||||
pub fn format_warning(msg: &str) -> String {
|
||||
format!("{} {}", "Warning:".yellow().bold(), msg)
|
||||
}
|
||||
|
||||
/// Format info message
|
||||
pub fn format_info(msg: &str) -> String {
|
||||
format!("{} {}", "ℹ".blue().bold(), msg)
|
||||
}
|
||||
|
||||
/// Export vector entries to JSON
|
||||
pub fn export_json(entries: &[VectorEntry]) -> anyhow::Result<String> {
|
||||
serde_json::to_string_pretty(entries)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to serialize to JSON: {}", e))
|
||||
}
|
||||
|
||||
/// Export vector entries to CSV
|
||||
pub fn export_csv(entries: &[VectorEntry]) -> anyhow::Result<String> {
|
||||
let mut wtr = csv::Writer::from_writer(vec![]);
|
||||
|
||||
// Write header
|
||||
wtr.write_record(&["id", "vector", "metadata"])?;
|
||||
|
||||
// Write entries
|
||||
for entry in entries {
|
||||
wtr.write_record(&[
|
||||
entry.id.as_ref().map(|s| s.as_str()).unwrap_or(""),
|
||||
&serde_json::to_string(&entry.vector)?,
|
||||
&serde_json::to_string(&entry.metadata)?,
|
||||
])?;
|
||||
}
|
||||
|
||||
wtr.flush()?;
|
||||
String::from_utf8(wtr.into_inner()?)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to convert CSV to string: {}", e))
|
||||
}
|
||||
|
||||
// Graph-specific formatting functions
|
||||
|
||||
/// Format graph node for display
|
||||
pub fn format_graph_node(
|
||||
id: &str,
|
||||
labels: &[String],
|
||||
properties: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> String {
|
||||
let mut output = String::new();
|
||||
|
||||
output.push_str(&format!("{} ({})\n", id.bold(), labels.join(":").cyan()));
|
||||
|
||||
if !properties.is_empty() {
|
||||
output.push_str(" Properties:\n");
|
||||
for (key, value) in properties {
|
||||
output.push_str(&format!(" {}: {}\n", key.yellow(), value));
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// Format graph relationship for display
|
||||
pub fn format_graph_relationship(
|
||||
id: &str,
|
||||
rel_type: &str,
|
||||
start_node: &str,
|
||||
end_node: &str,
|
||||
properties: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> String {
|
||||
let mut output = String::new();
|
||||
|
||||
output.push_str(&format!(
|
||||
"{} -[{}]-> {}\n",
|
||||
start_node.cyan(),
|
||||
rel_type.yellow(),
|
||||
end_node.cyan()
|
||||
));
|
||||
|
||||
if !properties.is_empty() {
|
||||
output.push_str(" Properties:\n");
|
||||
for (key, value) in properties {
|
||||
output.push_str(&format!(" {}: {}\n", key.yellow(), value));
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// Format graph query results as table
|
||||
pub fn format_graph_table(headers: &[String], rows: &[Vec<String>]) -> String {
|
||||
use prettytable::{Cell, Row, Table};
|
||||
|
||||
let mut table = Table::new();
|
||||
|
||||
// Add headers
|
||||
let header_cells: Vec<Cell> = headers
|
||||
.iter()
|
||||
.map(|h| Cell::new(h).style_spec("Fyb"))
|
||||
.collect();
|
||||
table.add_row(Row::new(header_cells));
|
||||
|
||||
// Add rows
|
||||
for row in rows {
|
||||
let cells: Vec<Cell> = row.iter().map(|v| Cell::new(v)).collect();
|
||||
table.add_row(Row::new(cells));
|
||||
}
|
||||
|
||||
table.to_string()
|
||||
}
|
||||
|
||||
/// Format graph statistics
|
||||
pub fn format_graph_stats(
|
||||
node_count: usize,
|
||||
rel_count: usize,
|
||||
label_count: usize,
|
||||
rel_type_count: usize,
|
||||
) -> String {
|
||||
format!(
|
||||
"\n{}\n Nodes: {}\n Relationships: {}\n Labels: {}\n Relationship Types: {}\n",
|
||||
"Graph Statistics".bold().green(),
|
||||
node_count.to_string().cyan(),
|
||||
rel_count.to_string().cyan(),
|
||||
label_count.to_string().cyan(),
|
||||
rel_type_count.to_string().cyan()
|
||||
)
|
||||
}
|
||||
552
vendor/ruvector/crates/ruvector-cli/src/cli/graph.rs
vendored
Normal file
552
vendor/ruvector/crates/ruvector-cli/src/cli/graph.rs
vendored
Normal file
@@ -0,0 +1,552 @@
|
||||
//! Graph database command implementations
|
||||
|
||||
use crate::cli::{format_error, format_info, format_success, ProgressTracker};
|
||||
use crate::config::Config;
|
||||
use anyhow::{Context, Result};
|
||||
use colored::*;
|
||||
use std::io::{self, BufRead, Write};
|
||||
use std::path::Path;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Graph database subcommands
|
||||
#[derive(clap::Subcommand, Debug)]
|
||||
pub enum GraphCommands {
|
||||
/// Create a new graph database
|
||||
Create {
|
||||
/// Database file path
|
||||
#[arg(short, long, default_value = "./ruvector-graph.db")]
|
||||
path: String,
|
||||
|
||||
/// Graph name
|
||||
#[arg(short, long, default_value = "default")]
|
||||
name: String,
|
||||
|
||||
/// Enable property indexing
|
||||
#[arg(long)]
|
||||
indexed: bool,
|
||||
},
|
||||
|
||||
/// Execute a Cypher query
|
||||
Query {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Cypher query to execute
|
||||
#[arg(short = 'q', long)]
|
||||
cypher: String,
|
||||
|
||||
/// Output format (table, json, csv)
|
||||
#[arg(long, default_value = "table")]
|
||||
format: String,
|
||||
|
||||
/// Show execution plan
|
||||
#[arg(long)]
|
||||
explain: bool,
|
||||
},
|
||||
|
||||
/// Interactive Cypher shell (REPL)
|
||||
Shell {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Enable multiline mode
|
||||
#[arg(long)]
|
||||
multiline: bool,
|
||||
},
|
||||
|
||||
/// Import data from file
|
||||
Import {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Input file path
|
||||
#[arg(short = 'i', long)]
|
||||
input: String,
|
||||
|
||||
/// Input format (csv, json, cypher)
|
||||
#[arg(long, default_value = "json")]
|
||||
format: String,
|
||||
|
||||
/// Graph name
|
||||
#[arg(short = 'g', long, default_value = "default")]
|
||||
graph: String,
|
||||
|
||||
/// Skip errors and continue
|
||||
#[arg(long)]
|
||||
skip_errors: bool,
|
||||
},
|
||||
|
||||
/// Export graph data to file
|
||||
Export {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Output file path
|
||||
#[arg(short = 'o', long)]
|
||||
output: String,
|
||||
|
||||
/// Output format (json, csv, cypher, graphml)
|
||||
#[arg(long, default_value = "json")]
|
||||
format: String,
|
||||
|
||||
/// Graph name
|
||||
#[arg(short = 'g', long, default_value = "default")]
|
||||
graph: String,
|
||||
},
|
||||
|
||||
/// Show graph database information
|
||||
Info {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Show detailed statistics
|
||||
#[arg(long)]
|
||||
detailed: bool,
|
||||
},
|
||||
|
||||
/// Run graph benchmarks
|
||||
Benchmark {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Number of queries to run
|
||||
#[arg(short = 'n', long, default_value = "1000")]
|
||||
queries: usize,
|
||||
|
||||
/// Benchmark type (traverse, pattern, aggregate)
|
||||
#[arg(short = 't', long, default_value = "traverse")]
|
||||
bench_type: String,
|
||||
},
|
||||
|
||||
/// Start HTTP/gRPC server
|
||||
Serve {
|
||||
/// Database file path
|
||||
#[arg(short = 'b', long, default_value = "./ruvector-graph.db")]
|
||||
db: String,
|
||||
|
||||
/// Server host
|
||||
#[arg(long, default_value = "127.0.0.1")]
|
||||
host: String,
|
||||
|
||||
/// HTTP port
|
||||
#[arg(long, default_value = "8080")]
|
||||
http_port: u16,
|
||||
|
||||
/// gRPC port
|
||||
#[arg(long, default_value = "50051")]
|
||||
grpc_port: u16,
|
||||
|
||||
/// Enable GraphQL endpoint
|
||||
#[arg(long)]
|
||||
graphql: bool,
|
||||
},
|
||||
}
|
||||
|
||||
/// Create a new graph database
|
||||
pub fn create_graph(path: &str, name: &str, indexed: bool, config: &Config) -> Result<()> {
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Creating graph database at: {}", path))
|
||||
);
|
||||
println!(" Graph name: {}", name.cyan());
|
||||
println!(
|
||||
" Property indexing: {}",
|
||||
if indexed {
|
||||
"enabled".green()
|
||||
} else {
|
||||
"disabled".dimmed()
|
||||
}
|
||||
);
|
||||
|
||||
// TODO: Integrate with ruvector-neo4j when available
|
||||
// For now, create a placeholder implementation
|
||||
std::fs::create_dir_all(Path::new(path).parent().unwrap_or(Path::new(".")))?;
|
||||
|
||||
println!("{}", format_success("Graph database created successfully!"));
|
||||
println!(
|
||||
"{}",
|
||||
format_info("Use 'ruvector graph shell' to start interactive mode")
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Execute a Cypher query
|
||||
pub fn execute_query(
|
||||
db_path: &str,
|
||||
cypher: &str,
|
||||
format: &str,
|
||||
explain: bool,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
if explain {
|
||||
println!("{}", "Query Execution Plan:".bold().cyan());
|
||||
println!("{}", format_info("EXPLAIN mode - showing query plan"));
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// TODO: Integrate with ruvector-neo4j Neo4jGraph implementation
|
||||
// Placeholder for actual query execution
|
||||
println!("{}", format_success("Executing Cypher query..."));
|
||||
println!(" Query: {}", cypher.dimmed());
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
match format {
|
||||
"table" => {
|
||||
println!("\n{}", format_graph_results_table(&[], cypher));
|
||||
}
|
||||
"json" => {
|
||||
println!("{}", format_graph_results_json(&[])?);
|
||||
}
|
||||
"csv" => {
|
||||
println!("{}", format_graph_results_csv(&[])?);
|
||||
}
|
||||
_ => return Err(anyhow::anyhow!("Unsupported output format: {}", format)),
|
||||
}
|
||||
|
||||
println!(
|
||||
"\n{}",
|
||||
format!("Query completed in {:.2}ms", elapsed.as_secs_f64() * 1000.0).dimmed()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Interactive Cypher shell (REPL)
|
||||
pub fn run_shell(db_path: &str, multiline: bool, config: &Config) -> Result<()> {
|
||||
println!("{}", "RuVector Graph Shell".bold().green());
|
||||
println!("Database: {}", db_path.cyan());
|
||||
println!(
|
||||
"Type {} to exit, {} for help\n",
|
||||
":exit".yellow(),
|
||||
":help".yellow()
|
||||
);
|
||||
|
||||
let stdin = io::stdin();
|
||||
let mut stdout = io::stdout();
|
||||
let mut query_buffer = String::new();
|
||||
|
||||
loop {
|
||||
// Print prompt
|
||||
if multiline && !query_buffer.is_empty() {
|
||||
print!("{}", " ... ".dimmed());
|
||||
} else {
|
||||
print!("{}", "cypher> ".green().bold());
|
||||
}
|
||||
stdout.flush()?;
|
||||
|
||||
// Read line
|
||||
let mut line = String::new();
|
||||
stdin.lock().read_line(&mut line)?;
|
||||
let line = line.trim();
|
||||
|
||||
// Handle special commands
|
||||
match line {
|
||||
":exit" | ":quit" | ":q" => {
|
||||
println!("{}", format_success("Goodbye!"));
|
||||
break;
|
||||
}
|
||||
":help" | ":h" => {
|
||||
print_shell_help();
|
||||
continue;
|
||||
}
|
||||
":clear" => {
|
||||
query_buffer.clear();
|
||||
println!("{}", format_info("Query buffer cleared"));
|
||||
continue;
|
||||
}
|
||||
"" => {
|
||||
if !multiline || query_buffer.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// In multiline mode, empty line executes query
|
||||
}
|
||||
_ => {
|
||||
query_buffer.push_str(line);
|
||||
query_buffer.push(' ');
|
||||
|
||||
if multiline && !line.ends_with(';') {
|
||||
continue; // Continue reading in multiline mode
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute query
|
||||
let query = query_buffer.trim().trim_end_matches(';');
|
||||
if !query.is_empty() {
|
||||
match execute_query(db_path, query, "table", false, config) {
|
||||
Ok(_) => {}
|
||||
Err(e) => println!("{}", format_error(&e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
query_buffer.clear();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Import graph data from file
|
||||
pub fn import_graph(
|
||||
db_path: &str,
|
||||
input_file: &str,
|
||||
format: &str,
|
||||
graph_name: &str,
|
||||
skip_errors: bool,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Importing graph data from: {}", input_file))
|
||||
);
|
||||
println!(" Format: {}", format.cyan());
|
||||
println!(" Graph: {}", graph_name.cyan());
|
||||
println!(
|
||||
" Skip errors: {}",
|
||||
if skip_errors {
|
||||
"yes".yellow()
|
||||
} else {
|
||||
"no".dimmed()
|
||||
}
|
||||
);
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// TODO: Implement actual import logic with ruvector-neo4j
|
||||
match format {
|
||||
"csv" => {
|
||||
println!("{}", format_info("Parsing CSV file..."));
|
||||
// Parse CSV and create nodes/relationships
|
||||
}
|
||||
"json" => {
|
||||
println!("{}", format_info("Parsing JSON file..."));
|
||||
// Parse JSON and create graph structure
|
||||
}
|
||||
"cypher" => {
|
||||
println!("{}", format_info("Executing Cypher statements..."));
|
||||
// Execute Cypher commands from file
|
||||
}
|
||||
_ => return Err(anyhow::anyhow!("Unsupported import format: {}", format)),
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!(
|
||||
"Import completed in {:.2}s",
|
||||
elapsed.as_secs_f64()
|
||||
))
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Export graph data to file
|
||||
pub fn export_graph(
|
||||
db_path: &str,
|
||||
output_file: &str,
|
||||
format: &str,
|
||||
graph_name: &str,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!("Exporting graph to: {}", output_file))
|
||||
);
|
||||
println!(" Format: {}", format.cyan());
|
||||
println!(" Graph: {}", graph_name.cyan());
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// TODO: Implement actual export logic with ruvector-neo4j
|
||||
match format {
|
||||
"json" => {
|
||||
println!("{}", format_info("Generating JSON export..."));
|
||||
// Export as JSON graph format
|
||||
}
|
||||
"csv" => {
|
||||
println!("{}", format_info("Generating CSV export..."));
|
||||
// Export nodes and edges as CSV files
|
||||
}
|
||||
"cypher" => {
|
||||
println!("{}", format_info("Generating Cypher statements..."));
|
||||
// Export as Cypher CREATE statements
|
||||
}
|
||||
"graphml" => {
|
||||
println!("{}", format_info("Generating GraphML export..."));
|
||||
// Export as GraphML XML format
|
||||
}
|
||||
_ => return Err(anyhow::anyhow!("Unsupported export format: {}", format)),
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
println!(
|
||||
"{}",
|
||||
format_success(&format!(
|
||||
"Export completed in {:.2}s",
|
||||
elapsed.as_secs_f64()
|
||||
))
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Show graph database information
|
||||
pub fn show_graph_info(db_path: &str, detailed: bool, config: &Config) -> Result<()> {
|
||||
println!("\n{}", "Graph Database Statistics".bold().green());
|
||||
|
||||
// TODO: Integrate with ruvector-neo4j to get actual statistics
|
||||
println!(" Database: {}", db_path.cyan());
|
||||
println!(" Graphs: {}", "1".cyan());
|
||||
println!(" Total nodes: {}", "0".cyan());
|
||||
println!(" Total relationships: {}", "0".cyan());
|
||||
println!(" Node labels: {}", "0".cyan());
|
||||
println!(" Relationship types: {}", "0".cyan());
|
||||
|
||||
if detailed {
|
||||
println!("\n{}", "Storage Information:".bold().cyan());
|
||||
println!(" Store size: {}", "0 bytes".cyan());
|
||||
println!(" Index size: {}", "0 bytes".cyan());
|
||||
|
||||
println!("\n{}", "Configuration:".bold().cyan());
|
||||
println!(" Cache size: {}", "N/A".cyan());
|
||||
println!(" Page size: {}", "N/A".cyan());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run graph benchmarks
|
||||
pub fn run_graph_benchmark(
|
||||
db_path: &str,
|
||||
num_queries: usize,
|
||||
bench_type: &str,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
println!("{}", "Running graph benchmark...".bold().green());
|
||||
println!(" Benchmark type: {}", bench_type.cyan());
|
||||
println!(" Queries: {}", num_queries.to_string().cyan());
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// TODO: Implement actual benchmarks with ruvector-neo4j
|
||||
match bench_type {
|
||||
"traverse" => {
|
||||
println!("{}", format_info("Benchmarking graph traversal..."));
|
||||
// Run traversal queries
|
||||
}
|
||||
"pattern" => {
|
||||
println!("{}", format_info("Benchmarking pattern matching..."));
|
||||
// Run pattern matching queries
|
||||
}
|
||||
"aggregate" => {
|
||||
println!("{}", format_info("Benchmarking aggregations..."));
|
||||
// Run aggregation queries
|
||||
}
|
||||
_ => return Err(anyhow::anyhow!("Unknown benchmark type: {}", bench_type)),
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
let qps = num_queries as f64 / elapsed.as_secs_f64();
|
||||
let avg_latency = elapsed.as_secs_f64() * 1000.0 / num_queries as f64;
|
||||
|
||||
println!("\n{}", "Benchmark Results:".bold().green());
|
||||
println!(" Total time: {:.2}s", elapsed.as_secs_f64());
|
||||
println!(" Queries per second: {:.0}", qps.to_string().cyan());
|
||||
println!(" Average latency: {:.2}ms", avg_latency.to_string().cyan());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Start HTTP/gRPC server
|
||||
pub fn serve_graph(
|
||||
db_path: &str,
|
||||
host: &str,
|
||||
http_port: u16,
|
||||
grpc_port: u16,
|
||||
enable_graphql: bool,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
println!("{}", "Starting RuVector Graph Server...".bold().green());
|
||||
println!(" Database: {}", db_path.cyan());
|
||||
println!(
|
||||
" HTTP endpoint: {}:{}",
|
||||
host.cyan(),
|
||||
http_port.to_string().cyan()
|
||||
);
|
||||
println!(
|
||||
" gRPC endpoint: {}:{}",
|
||||
host.cyan(),
|
||||
grpc_port.to_string().cyan()
|
||||
);
|
||||
|
||||
if enable_graphql {
|
||||
println!(
|
||||
" GraphQL endpoint: {}:{}/graphql",
|
||||
host.cyan(),
|
||||
http_port.to_string().cyan()
|
||||
);
|
||||
}
|
||||
|
||||
println!("\n{}", format_info("Server configuration loaded"));
|
||||
|
||||
// TODO: Implement actual server with ruvector-neo4j
|
||||
println!("{}", format_success("Server ready! Press Ctrl+C to stop."));
|
||||
|
||||
// Placeholder - would run actual server here
|
||||
println!(
|
||||
"\n{}",
|
||||
format_info("Server implementation pending - integrate with ruvector-neo4j")
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Helper functions for formatting graph results
|
||||
|
||||
fn format_graph_results_table(results: &[serde_json::Value], query: &str) -> String {
|
||||
let mut output = String::new();
|
||||
|
||||
if results.is_empty() {
|
||||
output.push_str(&format!("{}\n", "No results found".dimmed()));
|
||||
output.push_str(&format!("Query: {}\n", query.dimmed()));
|
||||
} else {
|
||||
output.push_str(&format!("{} results\n", results.len().to_string().cyan()));
|
||||
// TODO: Format results as table
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
fn format_graph_results_json(results: &[serde_json::Value]) -> Result<String> {
|
||||
serde_json::to_string_pretty(&results)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to serialize results: {}", e))
|
||||
}
|
||||
|
||||
fn format_graph_results_csv(results: &[serde_json::Value]) -> Result<String> {
|
||||
// TODO: Implement CSV formatting
|
||||
Ok(String::new())
|
||||
}
|
||||
|
||||
fn print_shell_help() {
|
||||
println!("\n{}", "RuVector Graph Shell Commands".bold().cyan());
|
||||
println!(" {} - Exit the shell", ":exit, :quit, :q".yellow());
|
||||
println!(
|
||||
" {} - Show this help message",
|
||||
":help, :h".yellow()
|
||||
);
|
||||
println!(" {} - Clear query buffer", ":clear".yellow());
|
||||
println!("\n{}", "Cypher Examples:".bold().cyan());
|
||||
println!(" {}", "CREATE (n:Person {{name: 'Alice'}})".dimmed());
|
||||
println!(" {}", "MATCH (n:Person) RETURN n".dimmed());
|
||||
println!(" {}", "MATCH (a)-[r:KNOWS]->(b) RETURN a, r, b".dimmed());
|
||||
println!();
|
||||
}
|
||||
2507
vendor/ruvector/crates/ruvector-cli/src/cli/hooks.rs
vendored
Normal file
2507
vendor/ruvector/crates/ruvector-cli/src/cli/hooks.rs
vendored
Normal file
File diff suppressed because it is too large
Load Diff
415
vendor/ruvector/crates/ruvector-cli/src/cli/hooks_postgres.rs
vendored
Normal file
415
vendor/ruvector/crates/ruvector-cli/src/cli/hooks_postgres.rs
vendored
Normal file
@@ -0,0 +1,415 @@
|
||||
//! PostgreSQL storage backend for hooks intelligence data
|
||||
//!
|
||||
//! This module provides PostgreSQL-based storage for the hooks system,
|
||||
//! using the ruvector extension for vector operations.
|
||||
//!
|
||||
//! Enable with the `postgres` feature flag.
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::{Config, Pool, Runtime};
|
||||
#[cfg(feature = "postgres")]
|
||||
use tokio_postgres::NoTls;
|
||||
|
||||
use std::env;
|
||||
|
||||
/// PostgreSQL storage configuration
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresConfig {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub user: String,
|
||||
pub password: Option<String>,
|
||||
pub dbname: String,
|
||||
}
|
||||
|
||||
impl PostgresConfig {
|
||||
/// Create config from environment variables
|
||||
pub fn from_env() -> Option<Self> {
|
||||
// Try RUVECTOR_POSTGRES_URL first, then DATABASE_URL
|
||||
if let Ok(url) = env::var("RUVECTOR_POSTGRES_URL").or_else(|_| env::var("DATABASE_URL")) {
|
||||
return Self::from_url(&url);
|
||||
}
|
||||
|
||||
// Try individual environment variables
|
||||
let host = env::var("RUVECTOR_PG_HOST").unwrap_or_else(|_| "localhost".to_string());
|
||||
let port = env::var("RUVECTOR_PG_PORT")
|
||||
.ok()
|
||||
.and_then(|p| p.parse().ok())
|
||||
.unwrap_or(5432);
|
||||
let user = env::var("RUVECTOR_PG_USER").ok()?;
|
||||
let password = env::var("RUVECTOR_PG_PASSWORD").ok();
|
||||
let dbname = env::var("RUVECTOR_PG_DATABASE").unwrap_or_else(|_| "ruvector".to_string());
|
||||
|
||||
Some(Self {
|
||||
host,
|
||||
port,
|
||||
user,
|
||||
password,
|
||||
dbname,
|
||||
})
|
||||
}
|
||||
|
||||
/// Parse PostgreSQL connection URL
|
||||
pub fn from_url(url: &str) -> Option<Self> {
|
||||
// Parse postgres://user:password@host:port/dbname
|
||||
let url = url
|
||||
.strip_prefix("postgres://")
|
||||
.or_else(|| url.strip_prefix("postgresql://"))?;
|
||||
|
||||
let (auth, rest) = url.split_once('@')?;
|
||||
let (user, password) = if auth.contains(':') {
|
||||
let (u, p) = auth.split_once(':')?;
|
||||
(u.to_string(), Some(p.to_string()))
|
||||
} else {
|
||||
(auth.to_string(), None)
|
||||
};
|
||||
|
||||
let (host_port, dbname) = rest.split_once('/')?;
|
||||
let dbname = dbname.split('?').next()?.to_string();
|
||||
|
||||
let (host, port) = if host_port.contains(':') {
|
||||
let (h, p) = host_port.split_once(':')?;
|
||||
(h.to_string(), p.parse().ok()?)
|
||||
} else {
|
||||
(host_port.to_string(), 5432)
|
||||
};
|
||||
|
||||
Some(Self {
|
||||
host,
|
||||
port,
|
||||
user,
|
||||
password,
|
||||
dbname,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// PostgreSQL storage backend for hooks
|
||||
#[cfg(feature = "postgres")]
|
||||
pub struct PostgresStorage {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl PostgresStorage {
|
||||
/// Create a new PostgreSQL storage backend
|
||||
pub async fn new(config: PostgresConfig) -> Result<Self, Box<dyn std::error::Error>> {
|
||||
let mut cfg = Config::new();
|
||||
cfg.host = Some(config.host);
|
||||
cfg.port = Some(config.port);
|
||||
cfg.user = Some(config.user);
|
||||
cfg.password = config.password;
|
||||
cfg.dbname = Some(config.dbname);
|
||||
|
||||
let pool = cfg.create_pool(Some(Runtime::Tokio1), NoTls)?;
|
||||
|
||||
Ok(Self { pool })
|
||||
}
|
||||
|
||||
/// Update Q-value for state-action pair
|
||||
pub async fn update_q(
|
||||
&self,
|
||||
state: &str,
|
||||
action: &str,
|
||||
reward: f32,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
client
|
||||
.execute(
|
||||
"SELECT ruvector_hooks_update_q($1, $2, $3)",
|
||||
&[&state, &action, &reward],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get best action for state
|
||||
pub async fn best_action(
|
||||
&self,
|
||||
state: &str,
|
||||
actions: &[String],
|
||||
) -> Result<Option<(String, f32, f32)>, Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
let row = client
|
||||
.query_opt(
|
||||
"SELECT action, q_value, confidence FROM ruvector_hooks_best_action($1, $2)",
|
||||
&[&state, &actions],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(row.map(|r| (r.get(0), r.get(1), r.get(2))))
|
||||
}
|
||||
|
||||
/// Store content in semantic memory
|
||||
pub async fn remember(
|
||||
&self,
|
||||
memory_type: &str,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<i32, Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
let metadata_str = serde_json::to_string(metadata)?;
|
||||
let row = client
|
||||
.query_one(
|
||||
"SELECT ruvector_hooks_remember($1, $2, $3, $4::jsonb)",
|
||||
&[&memory_type, &content, &embedding, &metadata_str],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(row.get(0))
|
||||
}
|
||||
|
||||
/// Search memory semantically
|
||||
pub async fn recall(
|
||||
&self,
|
||||
query_embedding: &[f32],
|
||||
limit: i32,
|
||||
) -> Result<Vec<MemoryResult>, Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
let rows = client
|
||||
.query(
|
||||
"SELECT id, memory_type, content, metadata::text, similarity
|
||||
FROM ruvector_hooks_recall($1, $2)",
|
||||
&[&query_embedding, &limit],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.map(|r| {
|
||||
let metadata_str: String = r.get(3);
|
||||
MemoryResult {
|
||||
id: r.get(0),
|
||||
memory_type: r.get(1),
|
||||
content: r.get(2),
|
||||
metadata: serde_json::from_str(&metadata_str).unwrap_or_default(),
|
||||
similarity: r.get(4),
|
||||
}
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Record file sequence
|
||||
pub async fn record_sequence(
|
||||
&self,
|
||||
from_file: &str,
|
||||
to_file: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
client
|
||||
.execute(
|
||||
"SELECT ruvector_hooks_record_sequence($1, $2)",
|
||||
&[&from_file, &to_file],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get suggested next files
|
||||
pub async fn suggest_next(
|
||||
&self,
|
||||
file: &str,
|
||||
limit: i32,
|
||||
) -> Result<Vec<(String, i32)>, Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
let rows = client
|
||||
.query(
|
||||
"SELECT to_file, count FROM ruvector_hooks_suggest_next($1, $2)",
|
||||
&[&file, &limit],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(rows.iter().map(|r| (r.get(0), r.get(1))).collect())
|
||||
}
|
||||
|
||||
/// Record error pattern
|
||||
pub async fn record_error(
|
||||
&self,
|
||||
code: &str,
|
||||
error_type: &str,
|
||||
message: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
client
|
||||
.execute(
|
||||
"SELECT ruvector_hooks_record_error($1, $2, $3)",
|
||||
&[&code, &error_type, &message],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Register agent in swarm
|
||||
pub async fn swarm_register(
|
||||
&self,
|
||||
agent_id: &str,
|
||||
agent_type: &str,
|
||||
capabilities: &[String],
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
client
|
||||
.execute(
|
||||
"SELECT ruvector_hooks_swarm_register($1, $2, $3)",
|
||||
&[&agent_id, &agent_type, &capabilities],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Record coordination between agents
|
||||
pub async fn swarm_coordinate(
|
||||
&self,
|
||||
source: &str,
|
||||
target: &str,
|
||||
weight: f32,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
client
|
||||
.execute(
|
||||
"SELECT ruvector_hooks_swarm_coordinate($1, $2, $3)",
|
||||
&[&source, &target, &weight],
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get swarm statistics
|
||||
pub async fn swarm_stats(&self) -> Result<SwarmStats, Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
let row = client
|
||||
.query_one("SELECT * FROM ruvector_hooks_swarm_stats()", &[])
|
||||
.await?;
|
||||
|
||||
Ok(SwarmStats {
|
||||
total_agents: row.get(0),
|
||||
active_agents: row.get(1),
|
||||
total_edges: row.get(2),
|
||||
avg_success_rate: row.get(3),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get overall statistics
|
||||
pub async fn get_stats(&self) -> Result<IntelligenceStats, Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
let row = client
|
||||
.query_one("SELECT * FROM ruvector_hooks_get_stats()", &[])
|
||||
.await?;
|
||||
|
||||
Ok(IntelligenceStats {
|
||||
total_patterns: row.get(0),
|
||||
total_memories: row.get(1),
|
||||
total_trajectories: row.get(2),
|
||||
total_errors: row.get(3),
|
||||
session_count: row.get(4),
|
||||
})
|
||||
}
|
||||
|
||||
/// Start a new session
|
||||
pub async fn session_start(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let client = self.pool.get().await?;
|
||||
client
|
||||
.execute("SELECT ruvector_hooks_session_start()", &[])
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Memory search result
|
||||
#[derive(Debug)]
|
||||
pub struct MemoryResult {
|
||||
pub id: i32,
|
||||
pub memory_type: String,
|
||||
pub content: String,
|
||||
pub metadata: serde_json::Value,
|
||||
pub similarity: f32,
|
||||
}
|
||||
|
||||
/// Swarm statistics
|
||||
#[derive(Debug)]
|
||||
pub struct SwarmStats {
|
||||
pub total_agents: i64,
|
||||
pub active_agents: i64,
|
||||
pub total_edges: i64,
|
||||
pub avg_success_rate: f32,
|
||||
}
|
||||
|
||||
/// Intelligence statistics
|
||||
#[derive(Debug)]
|
||||
pub struct IntelligenceStats {
|
||||
pub total_patterns: i64,
|
||||
pub total_memories: i64,
|
||||
pub total_trajectories: i64,
|
||||
pub total_errors: i64,
|
||||
pub session_count: i64,
|
||||
}
|
||||
|
||||
/// Check if PostgreSQL is available
|
||||
pub fn is_postgres_available() -> bool {
|
||||
PostgresConfig::from_env().is_some()
|
||||
}
|
||||
|
||||
/// Storage backend selector
|
||||
pub enum StorageBackend {
|
||||
#[cfg(feature = "postgres")]
|
||||
Postgres(PostgresStorage),
|
||||
Json(super::Intelligence),
|
||||
}
|
||||
|
||||
impl StorageBackend {
|
||||
/// Create storage backend from environment
|
||||
#[cfg(feature = "postgres")]
|
||||
pub async fn from_env() -> Result<Self, Box<dyn std::error::Error>> {
|
||||
if let Some(config) = PostgresConfig::from_env() {
|
||||
match PostgresStorage::new(config).await {
|
||||
Ok(pg) => return Ok(Self::Postgres(pg)),
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: PostgreSQL unavailable ({}), using JSON fallback",
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Self::Json(super::Intelligence::new(
|
||||
super::get_intelligence_path(),
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
pub fn from_env() -> Self {
|
||||
Self::Json(super::Intelligence::new(super::get_intelligence_path()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_config_from_url() {
|
||||
let config =
|
||||
PostgresConfig::from_url("postgres://user:pass@localhost:5432/ruvector").unwrap();
|
||||
assert_eq!(config.host, "localhost");
|
||||
assert_eq!(config.port, 5432);
|
||||
assert_eq!(config.user, "user");
|
||||
assert_eq!(config.password, Some("pass".to_string()));
|
||||
assert_eq!(config.dbname, "ruvector");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_from_url_no_password() {
|
||||
let config = PostgresConfig::from_url("postgres://user@localhost/ruvector").unwrap();
|
||||
assert_eq!(config.user, "user");
|
||||
assert_eq!(config.password, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_from_url_with_query() {
|
||||
let config = PostgresConfig::from_url(
|
||||
"postgres://user:pass@localhost:5432/ruvector?sslmode=require",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(config.dbname, "ruvector");
|
||||
}
|
||||
}
|
||||
15
vendor/ruvector/crates/ruvector-cli/src/cli/mod.rs
vendored
Normal file
15
vendor/ruvector/crates/ruvector-cli/src/cli/mod.rs
vendored
Normal file
@@ -0,0 +1,15 @@
|
||||
//! CLI module for Ruvector
|
||||
|
||||
pub mod commands;
|
||||
pub mod format;
|
||||
pub mod graph;
|
||||
pub mod hooks;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub mod hooks_postgres;
|
||||
pub mod progress;
|
||||
|
||||
pub use commands::*;
|
||||
pub use format::*;
|
||||
pub use graph::*;
|
||||
pub use hooks::*;
|
||||
pub use progress::ProgressTracker;
|
||||
56
vendor/ruvector/crates/ruvector-cli/src/cli/progress.rs
vendored
Normal file
56
vendor/ruvector/crates/ruvector-cli/src/cli/progress.rs
vendored
Normal file
@@ -0,0 +1,56 @@
|
||||
// ! Progress tracking for CLI operations
|
||||
|
||||
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
|
||||
use std::time::Duration;
|
||||
|
||||
/// Progress tracker for long-running operations
|
||||
pub struct ProgressTracker {
|
||||
multi: MultiProgress,
|
||||
}
|
||||
|
||||
impl ProgressTracker {
|
||||
/// Create a new progress tracker
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
multi: MultiProgress::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a progress bar for an operation
|
||||
pub fn create_bar(&self, total: u64, message: &str) -> ProgressBar {
|
||||
let pb = self.multi.add(ProgressBar::new(total));
|
||||
pb.set_style(
|
||||
ProgressStyle::default_bar()
|
||||
.template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} ({eta})")
|
||||
.unwrap()
|
||||
.progress_chars("#>-")
|
||||
);
|
||||
pb.set_message(message.to_string());
|
||||
pb.enable_steady_tick(Duration::from_millis(100));
|
||||
pb
|
||||
}
|
||||
|
||||
/// Create a spinner for indeterminate operations
|
||||
pub fn create_spinner(&self, message: &str) -> ProgressBar {
|
||||
let pb = self.multi.add(ProgressBar::new_spinner());
|
||||
pb.set_style(
|
||||
ProgressStyle::default_spinner()
|
||||
.template("{spinner:.green} {msg}")
|
||||
.unwrap(),
|
||||
);
|
||||
pb.set_message(message.to_string());
|
||||
pb.enable_steady_tick(Duration::from_millis(100));
|
||||
pb
|
||||
}
|
||||
|
||||
/// Finish all progress bars
|
||||
pub fn finish_all(&self) {
|
||||
// Progress bars auto-finish when dropped
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProgressTracker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user