Recommendation Engine (Rust)¶

:dep uni-db = { path = "../../../crates/uni" }
:dep tokio = { version = "1", features = ["full"] }
:dep serde_json = "1"

use uni_db::{Uni, DataType, IndexType, ScalarType, VectorMetric, VectorAlgo, VectorIndexCfg};
use std::collections::HashMap;
use serde_json::json;

// Helper macro to run async code in evcxr
macro_rules! run {
    ($e:expr) => {
        tokio::runtime::Runtime::new().unwrap().block_on($e)
    };
}

let db_path = "./recommendation_db";

// Clean up any existing database
if std::path::Path::new(db_path).exists() {
    std::fs::remove_dir_all(db_path).unwrap();
}

let db = run!(Uni::open(db_path).build()).unwrap();
println!("Opened database at {}", db_path);

run!(async {
    db.schema()
        .label("User")
            .property("name", DataType::String)
        .label("Book")
            .property("name",      DataType::String)
            .property("genre",     DataType::String)
            .property("embedding", DataType::Vector { dimensions: 4 })
            .index("embedding", IndexType::Vector(VectorIndexCfg {
                algorithm: VectorAlgo::Flat,
                metric: VectorMetric::L2,
            }))
        .edge_type("PURCHASED", &["User"], &["Book"])
        .apply()
        .await
}).unwrap();

println!("Schema created");

// 4D embeddings: [tech, fiction, history, science]
let books = vec![
    HashMap::from([
        ("name".to_string(),      json!("Clean Code")),
        ("genre".to_string(),     json!("tech")),
        ("embedding".to_string(), json!([0.95, 0.05, 0.0,  0.0 ])),
    ]),
    HashMap::from([
        ("name".to_string(),      json!("The Pragmatic Programmer")),
        ("genre".to_string(),     json!("tech")),
        ("embedding".to_string(), json!([0.90, 0.10, 0.0,  0.0 ])),
    ]),
    HashMap::from([
        ("name".to_string(),      json!("Designing Data-Intensive Apps")),
        ("genre".to_string(),     json!("tech")),
        ("embedding".to_string(), json!([0.85, 0.0,  0.0,  0.15])),
    ]),
    HashMap::from([
        ("name".to_string(),      json!("Dune")),
        ("genre".to_string(),     json!("fiction")),
        ("embedding".to_string(), json!([0.0,  0.95, 0.0,  0.05])),
    ]),
    HashMap::from([
        ("name".to_string(),      json!("Foundation")),
        ("genre".to_string(),     json!("fiction")),
        ("embedding".to_string(), json!([0.0,  0.85, 0.0,  0.15])),
    ]),
    HashMap::from([
        ("name".to_string(),      json!("Sapiens")),
        ("genre".to_string(),     json!("history")),
        ("embedding".to_string(), json!([0.0,  0.05, 0.7,  0.25])),
    ]),
];

let book_vids = run!(db.bulk_insert_vertices("Book", books)).unwrap();
let (clean_code, pragmatic, ddia, dune, foundation, sapiens) =
    (book_vids[0], book_vids[1], book_vids[2], book_vids[3], book_vids[4], book_vids[5]);

// 4 users
let users = vec![
    HashMap::from([("name".to_string(), json!("Alice"))]),
    HashMap::from([("name".to_string(), json!("Bob"))]),
    HashMap::from([("name".to_string(), json!("Carol"))]),
    HashMap::from([("name".to_string(), json!("Dave"))]),
];
let user_vids = run!(db.bulk_insert_vertices("User", users)).unwrap();
let (alice, bob, carol, dave) = (user_vids[0], user_vids[1], user_vids[2], user_vids[3]);

// Purchase history
run!(db.bulk_insert_edges("PURCHASED", vec![
    (alice, clean_code,  HashMap::new()),
    (alice, pragmatic,   HashMap::new()),
    (bob,   clean_code,  HashMap::new()),
    (bob,   dune,        HashMap::new()),
    (carol, pragmatic,   HashMap::new()),
    (carol, foundation,  HashMap::new()),
    (dave,  dune,        HashMap::new()),
    (dave,  foundation,  HashMap::new()),
    (dave,  sapiens,     HashMap::new()),
])).unwrap();

run!(db.flush()).unwrap();
println!("Data ingested");

let query_collab = r#"
    MATCH (alice:User {name: 'Alice'})-[:PURCHASED]->(b:Book)<-[:PURCHASED]-(other:User)
    WHERE other._vid <> alice._vid
    MATCH (other)-[:PURCHASED]->(rec:Book)
    WHERE NOT (alice)-[:PURCHASED]->(rec)
    RETURN rec.name AS recommendation, COUNT(DISTINCT other) AS buyers
    ORDER BY buyers DESC
"#;

let results = run!(db.query(query_collab)).unwrap();
println!("Collaborative recommendations for Alice:");
for row in &results.rows {
    println!("  {:?}", row);
}

let query_vec = r#"
    CALL uni.vector.query('Book', 'embedding', [0.95, 0.05, 0.0, 0.0], 3)
    YIELD node, distance
    RETURN node.name AS title, node.genre AS genre, distance
    ORDER BY distance
"#;

let results = run!(db.query(query_vec)).unwrap();
println!("Top 3 books semantically similar to tech query:");
for row in &results.rows {
    println!("  {:?}", row);
}
// All 3 results should be tech books
assert!(results.rows.len() == 3, "Expected 3 results, got {}", results.rows.len());

let query_hybrid = r#"
    CALL uni.vector.query('Book', 'embedding', [0.0, 0.95, 0.0, 0.05], 3)
    YIELD node, distance
    MATCH (u:User)-[:PURCHASED]->(node)
    RETURN node.name AS book, u.name AS buyer, distance
    ORDER BY distance, buyer
"#;

let results = run!(db.query(query_hybrid)).unwrap();
println!("Fiction book buyers (via vector + graph):");
for row in &results.rows {
    println!("  {:?}", row);
}

let query_discovery = r#"
    MATCH (alice:User {name: 'Alice'})
    MATCH (u:User)-[:PURCHASED]->(b:Book)
    WHERE NOT (alice)-[:PURCHASED]->(b) AND u._vid <> alice._vid
    RETURN b.name AS book, COUNT(DISTINCT u) AS buyers
    ORDER BY buyers DESC
"#;

let results = run!(db.query(query_discovery)).unwrap();
println!("Popular books Alice has not read:");
for row in &results.rows {
    println!("  {:?}", row);
}