use agentdb::{AgentDB, BatchEntry, DistanceMetric, HybridQuery, SearchOptions, VectorEntry};
use serde_json::json;
fn embed(seed: f32, dim: usize) -> Vec<f32> {
(0..dim)
.map(|i| ((seed + i as f32) * 0.1).sin().abs())
.collect()
}
fn main() -> agentdb::Result<()> {
let db = AgentDB::open(":memory:");
let db = db?;
println!("=== AgentDB v0.2.0 — Query Power Demo ===\n");
println!("1. Batch upsert (single transaction)");
let col = db.vectors().collection("docs", 8)?;
let batch: Vec<BatchEntry> = vec![
BatchEntry {
id: "doc_rust".into(),
vector: embed(1.0, 8),
metadata: Some(
json!({"text": "Rust systems programming language", "lang": "en", "score": 9, "ts": 1700000100}),
),
},
BatchEntry {
id: "doc_agents".into(),
vector: embed(2.0, 8),
metadata: Some(
json!({"text": "AI agents and autonomous systems", "lang": "en", "score": 8, "ts": 1700000200}),
),
},
BatchEntry {
id: "doc_db".into(),
vector: embed(3.0, 8),
metadata: Some(
json!({"text": "Database design and vector search", "lang": "en", "score": 7, "ts": 1700000300}),
),
},
BatchEntry {
id: "doc_rag".into(),
vector: embed(4.0, 8),
metadata: Some(
json!({"text": "Retrieval augmented generation RAG", "lang": "en", "score": 8, "ts": 1700000400}),
),
},
BatchEntry {
id: "doc_memory".into(),
vector: embed(5.0, 8),
metadata: Some(
json!({"text": "Memory graphs for episodic recall", "lang": "en", "score": 9, "ts": 1700000500}),
),
},
BatchEntry {
id: "doc_fr".into(),
vector: embed(6.0, 8),
metadata: Some(
json!({"text": "Apprentissage automatique avancé", "lang": "fr", "score": 6, "ts": 1700000600}),
),
},
];
let inserted = col.upsert_batch(batch)?;
println!(" Inserted {} documents in one transaction", inserted);
println!(" Collection size: {}", col.count()?);
println!("\n2. Advanced metadata filtering");
let high_score = col.search(
&embed(1.0, 8),
SearchOptions {
top_k: 10,
metric: DistanceMetric::Cosine,
filter: Some(json!({ "score": { "$gt": 7 } })),
},
)?;
println!(" score > 7: {} results", high_score.len());
for r in &high_score {
println!(
" {} score={}",
r.id,
r.metadata.as_ref().unwrap()["score"]
);
}
let english = col.search(
&embed(1.0, 8),
SearchOptions {
top_k: 10,
metric: DistanceMetric::Cosine,
filter: Some(json!({ "lang": { "$in": ["en"] } })),
},
)?;
println!(" lang $in [en]: {} results", english.len());
let has_score = col.search(
&embed(1.0, 8),
SearchOptions {
top_k: 10,
metric: DistanceMetric::Cosine,
filter: Some(json!({ "score": { "$exists": true } })),
},
)?;
println!(" score $exists true: {} results", has_score.len());
let top_en = col.search(
&embed(1.0, 8),
SearchOptions {
top_k: 10,
metric: DistanceMetric::Cosine,
filter: Some(json!({ "lang": "en", "score": { "$gte": 8 } })),
},
)?;
println!(" lang=en AND score >= 8: {} results", top_en.len());
println!("\n3. Full-text search (FTS5 + BM25)");
let fts = db.fts();
let docs_text = vec![
(
"doc_rust",
"Rust systems programming language memory safety performance",
),
(
"doc_agents",
"AI agents autonomous systems planning reasoning memory",
),
(
"doc_db",
"database design vector search embeddings storage retrieval",
),
(
"doc_rag",
"retrieval augmented generation language model context",
),
(
"doc_memory",
"memory graphs episodic recall knowledge representation",
),
(
"doc_fr",
"apprentissage automatique intelligence artificielle",
),
];
let col_id = col.id.clone();
for (id, text) in &docs_text {
fts.index_text("docs", id, &col_id, text)?;
}
fts.optimize("docs")?;
let kw_results = fts.search("docs", "memory", 5)?;
println!(" FTS 'memory': {} results", kw_results.len());
for r in &kw_results {
println!(" {} | snippet: {}", r.id, r.snippet);
}
let kw2 = fts.search("docs", "vector search", 5)?;
println!(" FTS 'vector search': {} results", kw2.len());
for r in &kw2 {
println!(" {} | rank: {:.4}", r.id, r.rank);
}
println!("\n4. Hybrid query (graph traversal + vector search)");
let graph = db.memory();
graph.add_node("session_x", "session", Some(json!({"user": "harshal"})))?;
graph.add_node("doc_rust", "doc", Some(json!({"title": "Rust"})))?;
graph.add_node("doc_agents", "doc", Some(json!({"title": "Agents"})))?;
graph.add_node("doc_memory", "doc", Some(json!({"title": "Memory"})))?;
graph.add_edge("session_x", "doc_rust", "read", 0.95)?;
graph.add_edge("session_x", "doc_agents", "read", 0.80)?;
graph.add_edge("session_x", "doc_memory", "read", 0.70)?;
let hybrid_results = db.hybrid_query(HybridQuery {
anchor_node: "session_x",
embedding: &embed(1.0, 8),
collection: "docs",
graph_depth: 1,
top_k: 3,
alpha: 0.6,
filter: None,
})?;
println!(
" Top {} hybrid results (alpha=0.6):",
hybrid_results.len()
);
for r in &hybrid_results {
println!(
" {} | rank={:.4} vec={:.4} graph={:.2}",
r.id, r.rank_score, r.vector_score, r.graph_weight
);
}
let graph_only = db.hybrid_query(HybridQuery {
anchor_node: "session_x",
embedding: &embed(1.0, 8),
collection: "docs",
graph_depth: 1,
top_k: 3,
alpha: 0.0,
filter: None,
})?;
println!(
" Top {} results (alpha=0.0, pure graph):",
graph_only.len()
);
for r in &graph_only {
println!(" {} | graph_weight={:.2}", r.id, r.graph_weight);
}
let _ = VectorEntry {
id: "test".into(),
vector: vec![0.0; 8],
metadata: None,
};
println!("\n✓ v0.2.0 demo complete.");
Ok(())
}