swiftide-integrations 0.33.0

Composable LLM agents, typed task graphs, and streaming RAG pipelines in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

use async_trait::async_trait;
use swiftide_core::chat_completion::errors::LanguageModelError;
use swiftide_core::{SparseEmbedding, SparseEmbeddingModel, SparseEmbeddings};

use super::{EmbeddingModelType, FastEmbed};
#[async_trait]
impl SparseEmbeddingModel for FastEmbed {
    #[tracing::instrument(skip_all)]
    async fn sparse_embed(
        &self,
        input: Vec<String>,
    ) -> Result<SparseEmbeddings, LanguageModelError> {
        let mut embedding_model = self.embedding_model.lock().await;

        match &mut *embedding_model {
            EmbeddingModelType::Sparse(model) => model
                .embed(input, self.batch_size)
                .map_err(LanguageModelError::permanent)
                .and_then(|embeddings| {
                    embeddings
                        .into_iter()
                        .map(|embedding| {
                            let indices = embedding
                                .indices
                                .iter()
                                .map(|v| u32::try_from(*v).map_err(LanguageModelError::permanent))
                                .collect::<Result<Vec<_>, LanguageModelError>>()?;

                            Ok(SparseEmbedding {
                                indices,
                                values: embedding.values,
                            })
                        })
                        .collect()
                }),
            EmbeddingModelType::Dense(_) => Err(LanguageModelError::PermanentError(
                "Expected sparse model, got dense".into(),
            )),
        }
    }
}