Retrieval-Augmented Generation (RAG) with Uni (Rust)¶

Combining vector search with knowledge graph traversal for hybrid retrieval over Python web framework documentation.

: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 = "./rag_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);

1. Schema¶

Chunks of text with embeddings, linked to named Entities via MENTIONS edges.

run!(async {
    db.schema()
        .label("Chunk")
            .property("chunk_id",  DataType::String)
            .property("text",      DataType::String)
            .property("embedding", DataType::Vector { dimensions: 4 })
            .index("embedding", IndexType::Vector(VectorIndexCfg {
                algorithm: VectorAlgo::Flat,
                metric: VectorMetric::L2,
            }))
        .label("Entity")
            .property("name", DataType::String)
            .property("type", DataType::String)
        .edge_type("MENTIONS", &["Chunk"], &["Entity"])
        .apply()
        .await
}).unwrap();

println!("RAG schema created");

2. Ingest Data¶

// 4D embeddings: [auth, routing, database, testing]
let chunks = vec![
    HashMap::from([("chunk_id".to_string(), json!("c1")), ("text".to_string(), json!("JWT tokens issued by /auth/login endpoint. Tokens expire after 1 hour.")),        ("embedding".to_string(), json!([1.0,  0.0,  0.0,  0.0 ]))]),
    HashMap::from([("chunk_id".to_string(), json!("c2")), ("text".to_string(), json!("Token refresh via /auth/refresh. Send expired token, receive new one.")),         ("embedding".to_string(), json!([0.95, 0.05, 0.0,  0.0 ]))]),
    HashMap::from([("chunk_id".to_string(), json!("c3")), ("text".to_string(), json!("Password hashing uses bcrypt with cost factor 12.")),                             ("embedding".to_string(), json!([0.85, 0.0,  0.0,  0.15]))]),
    HashMap::from([("chunk_id".to_string(), json!("c4")), ("text".to_string(), json!("Routes defined with @app.route decorator. Supports GET, POST, PUT, DELETE.")),   ("embedding".to_string(), json!([0.0,  1.0,  0.0,  0.0 ]))]),
    HashMap::from([("chunk_id".to_string(), json!("c5")), ("text".to_string(), json!("Middleware intercepts requests before handlers. Register with app.use().")),      ("embedding".to_string(), json!([0.05, 0.9,  0.05, 0.0 ]))]),
    HashMap::from([("chunk_id".to_string(), json!("c6")), ("text".to_string(), json!("ConnectionPool manages DB connections. Max pool size defaults to 10.")),          ("embedding".to_string(), json!([0.0,  0.0,  1.0,  0.0 ]))]),
    HashMap::from([("chunk_id".to_string(), json!("c7")), ("text".to_string(), json!("ORM models inherit from BaseModel. Columns map to database fields.")),            ("embedding".to_string(), json!([0.0,  0.1,  0.9,  0.0 ]))]),
    HashMap::from([("chunk_id".to_string(), json!("c8")), ("text".to_string(), json!("TestClient simulates HTTP requests without starting a server.")),                 ("embedding".to_string(), json!([0.0,  0.2,  0.0,  0.8 ]))]),
];

let chunk_vids = run!(db.bulk_insert_vertices("Chunk", chunks)).unwrap();
let (c1, c2, c3, c4, c5, c6, c7, c8) =
    (chunk_vids[0], chunk_vids[1], chunk_vids[2], chunk_vids[3],
     chunk_vids[4], chunk_vids[5], chunk_vids[6], chunk_vids[7]);

// 6 entities
let entities = vec![
    HashMap::from([("name".to_string(), json!("JWT")),            ("type".to_string(), json!("technology"))]),
    HashMap::from([("name".to_string(), json!("authentication")), ("type".to_string(), json!("concept"))]),
    HashMap::from([("name".to_string(), json!("routing")),        ("type".to_string(), json!("concept"))]),
    HashMap::from([("name".to_string(), json!("database")),       ("type".to_string(), json!("concept"))]),
    HashMap::from([("name".to_string(), json!("bcrypt")),         ("type".to_string(), json!("technology"))]),
    HashMap::from([("name".to_string(), json!("ConnectionPool")), ("type".to_string(), json!("class"))]),
];

let entity_vids = run!(db.bulk_insert_vertices("Entity", entities)).unwrap();
let (jwt, auth_entity, routing_entity, db_entity, bcrypt_entity, pool_entity) =
    (entity_vids[0], entity_vids[1], entity_vids[2], entity_vids[3], entity_vids[4], entity_vids[5]);

// MENTIONS edges
run!(db.bulk_insert_edges("MENTIONS", vec![
    (c1, jwt,            HashMap::new()),
    (c1, auth_entity,    HashMap::new()),
    (c2, jwt,            HashMap::new()),
    (c2, auth_entity,    HashMap::new()),
    (c3, bcrypt_entity,  HashMap::new()),
    (c3, auth_entity,    HashMap::new()),
    (c4, routing_entity, HashMap::new()),
    (c5, routing_entity, HashMap::new()),
    (c6, db_entity,      HashMap::new()),
    (c6, pool_entity,    HashMap::new()),
    (c7, db_entity,      HashMap::new()),
])).unwrap();

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

3. Pure Vector Search¶

Find the 3 chunks most similar to an authentication query.

let query_vec = r#"
    CALL uni.vector.query('Chunk', 'embedding', [1.0, 0.0, 0.0, 0.0], 3)
    YIELD node, distance
    RETURN node.chunk_id AS chunk_id, node.text AS text, distance
    ORDER BY distance
"#;

let results = run!(db.query(query_vec)).unwrap();
println!("Top 3 chunks for auth query:");
for row in &results.rows {
    println!("  {:?}", row);
}
// Expected: c1, c2, c3 (auth chunks)
assert!(results.rows.len() == 3, "Expected 3 results, got {}", results.rows.len());

4. Graph Expansion¶

Same vector seeds — also show which entities each chunk mentions.

let query_expand = r#"
    CALL uni.vector.query('Chunk', 'embedding', [1.0, 0.0, 0.0, 0.0], 3)
    YIELD node, distance
    MATCH (node)-[:MENTIONS]->(e:Entity)
    RETURN node.chunk_id AS chunk_id, e.name AS entity, distance
    ORDER BY distance, entity
"#;

let results = run!(db.query(query_expand)).unwrap();
println!("Entities mentioned by top auth chunks:");
for row in &results.rows {
    println!("  {:?}", row);
}

5. Entity Bridging¶

Find all chunks related to the auth seeds via shared entity mentions — the core graph RAG technique.

let query_bridge = r#"
    CALL uni.vector.query('Chunk', 'embedding', [1.0, 0.0, 0.0, 0.0], 3)
    YIELD node AS anchor, distance
    MATCH (anchor)-[:MENTIONS]->(e:Entity)<-[:MENTIONS]-(related:Chunk)
    WHERE related._vid <> anchor._vid
    RETURN anchor.chunk_id AS anchor_id, e.name AS bridge_entity,
           related.chunk_id AS related_id
    ORDER BY anchor_id, bridge_entity
"#;

let results = run!(db.query(query_bridge)).unwrap();
println!("Entity bridges between auth chunks:");
for row in &results.rows {
    println!("  {:?}", row);
}

6. Context Assembly¶

Full hybrid pipeline: vector seeds + graph bridging → collect unique chunks for the LLM context window.

use std::collections::HashSet;

let query_ctx = r#"
    CALL uni.vector.query('Chunk', 'embedding', [1.0, 0.0, 0.0, 0.0], 3)
    YIELD node AS seed, distance
    MATCH (seed)-[:MENTIONS]->(e:Entity)<-[:MENTIONS]-(related:Chunk)
    RETURN seed.chunk_id AS seed_id, seed.text AS seed_text,
           related.chunk_id AS related_id, related.text AS related_text,
           e.name AS shared_entity
    ORDER BY seed_id, shared_entity
"#;

let results = run!(db.query(query_ctx)).unwrap();
println!("Context assembly result:");
println!("  {} rows retrieved for LLM context window", results.rows.len());
for row in &results.rows {
    println!("  {:?}", row);
}