kalosm_language_model/embedding/
into_embedding.rs

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
use futures_util::Future;

use crate::{Embedder, EmbedderExt, Embedding};

/// Convert a type into an embedding with an embedding model.
pub trait IntoEmbedding {
    /// Convert the type into an embedding with the given embedding model.
    fn into_embedding<E: Embedder>(
        self,
        embedder: &E,
    ) -> impl Future<Output = Result<Embedding, E::Error>> + Send;

    /// Convert the type into a query embedding with the given embedding model.
    fn into_query_embedding<E: Embedder>(
        self,
        embedder: &E,
    ) -> impl Future<Output = Result<Embedding, E::Error>> + Send;
}

/// Convert any type that implements [`ToString`] into an embedding with an embedding model.
impl<S: ToString + Send> IntoEmbedding for S {
    async fn into_embedding<E: Embedder>(self, embedder: &E) -> Result<Embedding, E::Error> {
        embedder.embed(self).await
    }

    async fn into_query_embedding<E: Embedder>(self, embedder: &E) -> Result<Embedding, E::Error> {
        embedder.embed_query(self).await
    }
}

/// Convert an embedding of the same vector space into an embedding with an embedding model.
impl IntoEmbedding for Embedding {
    async fn into_embedding<E: Embedder>(self, _: &E) -> Result<Embedding, E::Error> {
        Ok(self)
    }

    async fn into_query_embedding<E: Embedder>(self, _: &E) -> Result<Embedding, E::Error> {
        Ok(self)
    }
}