variphx's picture
add/fix: add search functionalities and refactor some modules
6b23458
raw
history blame contribute delete
751 Bytes
use tonic::Request;
use crate::models::protos::text_embed::{
TextRequest, text_embedding_service_client::TextEmbeddingServiceClient,
};
pub struct TextEmbeddingService {
client: TextEmbeddingServiceClient<tonic::transport::Channel>,
}
impl TextEmbeddingService {
pub async fn new() -> anyhow::Result<Self> {
Ok(Self {
client: TextEmbeddingServiceClient::connect(std::env::var("EMBEDDING_SERVER_URL")?)
.await?,
})
}
pub async fn embed_text(&mut self, text: &str) -> anyhow::Result<Vec<f32>> {
let request = Request::new(TextRequest { text: text.into() });
let response = self.client.get_embedding(request).await?;
Ok(response.into_inner().embeddings)
}
}