45 lines
1.4 KiB
Rust
45 lines
1.4 KiB
Rust
|
use std::{sync::Arc, time::Duration};
|
||
|
|
|
||
|
|
use stratum_embeddings::{
|
||
|
|
EmbeddingOptions, EmbeddingService, FastEmbedProvider, MemoryCache, OllamaProvider,
|
||
|
|
};
|
||
|
|
use tracing::{info, warn};
|
||
|
|
|
||
|
|
#[tokio::main]
|
||
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||
|
|
tracing_subscriber::fmt::init();
|
||
|
|
|
||
|
|
info!("Setting up primary provider (Ollama)...");
|
||
|
|
let primary = OllamaProvider::default_model()?;
|
||
|
|
|
||
|
|
info!("Setting up fallback provider (FastEmbed)...");
|
||
|
|
let fallback =
|
||
|
|
Arc::new(FastEmbedProvider::small()?) as Arc<dyn stratum_embeddings::EmbeddingProvider>;
|
||
|
|
|
||
|
|
let cache = MemoryCache::new(1000, Duration::from_secs(300));
|
||
|
|
let service = EmbeddingService::new(primary)
|
||
|
|
.with_cache(cache)
|
||
|
|
.with_fallback(fallback);
|
||
|
|
|
||
|
|
let options = EmbeddingOptions::default_with_cache();
|
||
|
|
|
||
|
|
info!("Checking if Ollama is available...");
|
||
|
|
if service.is_ready().await {
|
||
|
|
info!("Ollama is available, using as primary");
|
||
|
|
} else {
|
||
|
|
warn!("Ollama not available, will fall back to FastEmbed");
|
||
|
|
}
|
||
|
|
|
||
|
|
info!("Embedding text (will use available provider)...");
|
||
|
|
let text = "This demonstrates fallback strategy in action";
|
||
|
|
let embedding = service.embed(text, &options).await?;
|
||
|
|
|
||
|
|
info!(
|
||
|
|
"Successfully generated embedding with {} dimensions",
|
||
|
|
embedding.len()
|
||
|
|
);
|
||
|
|
info!("Cache size: {}", service.cache_size());
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|