use anyhow::Context;
use indicatif::ProgressIterator;
use llm::{
LLMProvider,
builder::{LLMBackend, LLMBuilder},
};
use serde::Serialize;
use crate::intake::dataset::DataSet;
pub struct EmbeddingProvider {
pub model: Box<dyn LLMProvider>,
}
impl EmbeddingProvider {
pub fn new(model: &str) -> anyhow::Result<Self> {
let base_url = std::env::var("OLLAMA_URL").unwrap_or("http://127.0.0.1:11434".into());
let llm = LLMBuilder::new()
.backend(LLMBackend::Ollama)
.base_url(base_url)
.model(model)
.build()
.context("Error creando modelo embdding")?;
Ok(EmbeddingProvider { model: llm })
}
pub fn new_openai(model: &str) -> anyhow::Result<Self> {
dotenvy::dotenv().context(".env absent")?;
let api_key = std::env::var("OPENAI_API_KEY").context("Api key absent")?;
let llm = LLMBuilder::new()
.backend(LLMBackend::OpenAI)
.api_key(api_key)
.model(model)
.build()
.context("Error creando modelo embdding")?;
Ok(EmbeddingProvider { model: llm })
}
pub async fn embed_properties<T>(&self, dataset: DataSet<T>) -> anyhow::Result<Vec<Vec<f32>>>
where
T: Serialize + Clone,
{
let mut embeddings = vec![];
let mut failed_ids = vec![];
for (n, chunk) in dataset
.data
.context("There is no data to embed")?
.chunks(50)
.enumerate()
.progress()
{
let properties_string_chunk: Vec<_> = chunk
.iter()
.map(|article| serde_json::to_string(article).expect("Couldn't serialize article"))
.collect();
let embeddings_chunk = self.model.embed(properties_string_chunk.clone()).await;
println!("{:?}", embeddings_chunk);
match embeddings_chunk {
Ok(emb) => embeddings.extend(emb),
Err(_) => failed_ids.push(n),
}
}
println!(
"Returning from embedding function. Failed chunk ids:\n{:#?}",
failed_ids
);
Ok(embeddings)
}
}