67 lines
2.3 KiB
Rust
67 lines
2.3 KiB
Rust
use std::time::Duration;
|
|
|
|
use stratum_embeddings::{
|
|
EmbeddingOptions, EmbeddingService, FastEmbedProvider, MemoryCache, SurrealDbStore,
|
|
VectorStore, VectorStoreConfig,
|
|
};
|
|
use tracing::info;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
tracing_subscriber::fmt::init();
|
|
|
|
info!("Initializing embedding service...");
|
|
let provider = FastEmbedProvider::small()?;
|
|
let cache = MemoryCache::new(1000, Duration::from_secs(300));
|
|
let service = EmbeddingService::new(provider).with_cache(cache);
|
|
|
|
info!("Creating SurrealDB in-memory store...");
|
|
let config = VectorStoreConfig::new(384);
|
|
let store = SurrealDbStore::new_memory("concepts", config).await?;
|
|
|
|
let concepts = vec![
|
|
("ownership", "Rust's ownership system prevents memory leaks"),
|
|
(
|
|
"borrowing",
|
|
"Borrowing allows references without ownership transfer",
|
|
),
|
|
("lifetimes", "Lifetimes ensure references remain valid"),
|
|
("traits", "Traits define shared behavior across types"),
|
|
("generics", "Generics enable code reuse with type safety"),
|
|
];
|
|
|
|
info!("Embedding and storing {} concepts...", concepts.len());
|
|
let options = EmbeddingOptions::default_with_cache();
|
|
|
|
for (id, description) in &concepts {
|
|
let embedding = service.embed(description, &options).await?;
|
|
let metadata = serde_json::json!({
|
|
"concept": id,
|
|
"description": description,
|
|
"language": "rust"
|
|
});
|
|
store.upsert(id, &embedding, metadata).await?;
|
|
}
|
|
|
|
info!("Concepts stored successfully");
|
|
|
|
info!("Performing knowledge graph search...");
|
|
let query = "How does Rust manage memory?";
|
|
let query_embedding = service.embed(query, &options).await?;
|
|
|
|
let results = store.search(&query_embedding, 3, None).await?;
|
|
|
|
info!("Most relevant concepts for: '{}'", query);
|
|
for (i, result) in results.iter().enumerate() {
|
|
let concept = result.metadata["concept"].as_str().unwrap_or("N/A");
|
|
let description = result.metadata["description"].as_str().unwrap_or("N/A");
|
|
info!(" {}. {} [score: {:.4}]", i + 1, concept, result.score);
|
|
info!(" {}", description);
|
|
}
|
|
|
|
let count = store.count().await?;
|
|
info!("Total concepts in graph: {}", count);
|
|
|
|
Ok(())
|
|
}
|