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

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

In [ ]:

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

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

In [ ]:

use uni::{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)
    };
}

use uni::{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) }; }

In [ ]:

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);

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.

In [ ]:

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");

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¶

In [ ]:

// 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");

// 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.

In [ ]:

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());

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.

In [ ]:

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);
}

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.

In [ ]:

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);
}

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.

In [ ]:

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);
}

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); }