64 lines
1.9 KiB
Rust
64 lines
1.9 KiB
Rust
|
// Test RLMEngine BM25 integration
|
||
|
|
use std::sync::Arc;
|
||
|
|
|
||
|
|
use surrealdb::engine::remote::ws::Ws;
|
||
|
|
use surrealdb::opt::auth::Root;
|
||
|
|
use surrealdb::Surreal;
|
||
|
|
use vapora_rlm::search::bm25::BM25Index;
|
||
|
|
use vapora_rlm::storage::SurrealDBStorage;
|
||
|
|
use vapora_rlm::RLMEngine;
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
#[ignore] // Requires SurrealDB
|
||
|
|
async fn test_engine_bm25_query() {
|
||
|
|
// Setup - same as E2E test
|
||
|
|
let db = Surreal::new::<Ws>("127.0.0.1:8000").await.unwrap();
|
||
|
|
db.signin(Root {
|
||
|
|
username: "root",
|
||
|
|
password: "root",
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.unwrap();
|
||
|
|
db.use_ns("test_engine_bm25")
|
||
|
|
.use_db("test_engine_bm25")
|
||
|
|
.await
|
||
|
|
.unwrap();
|
||
|
|
|
||
|
|
let storage = Arc::new(SurrealDBStorage::new(db));
|
||
|
|
let bm25_index = Arc::new(BM25Index::new().unwrap());
|
||
|
|
let engine = Arc::new(RLMEngine::new(storage, bm25_index).unwrap());
|
||
|
|
|
||
|
|
// Load a document
|
||
|
|
let doc_id = format!("test-{}", uuid::Uuid::new_v4());
|
||
|
|
let content = "This is test content with error handling patterns in Rust programming.";
|
||
|
|
|
||
|
|
println!("Loading document...");
|
||
|
|
let chunk_count = engine.load_document(&doc_id, content, None).await.unwrap();
|
||
|
|
println!("✓ Loaded {} chunks", chunk_count);
|
||
|
|
|
||
|
|
// Small delay to ensure commit completes
|
||
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||
|
|
|
||
|
|
// Query
|
||
|
|
println!("Querying for 'error handling'...");
|
||
|
|
let results = engine
|
||
|
|
.query(&doc_id, "error handling", None, 5)
|
||
|
|
.await
|
||
|
|
.unwrap();
|
||
|
|
println!("✓ Found {} results", results.len());
|
||
|
|
|
||
|
|
for (i, result) in results.iter().enumerate() {
|
||
|
|
println!(
|
||
|
|
" Result {}: score={}, content_preview={}",
|
||
|
|
i + 1,
|
||
|
|
result.score,
|
||
|
|
&result.chunk.content[..50.min(result.chunk.content.len())]
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
assert!(
|
||
|
|
!results.is_empty(),
|
||
|
|
"Should find results for 'error handling'"
|
||
|
|
);
|
||
|
|
}
|