Skip to main content

headless_lms_utils/
azure_embedding.rs

1use crate::http::REQWEST_CLIENT;
2use crate::prelude::*;
3use headless_lms_base::config::ApplicationConfiguration;
4use secrecy::ExposeSecret;
5
6use serde::{Deserialize, Serialize};
7
8#[derive(Serialize, Deserialize)]
9pub struct EmbeddingRequest {
10    pub model: String,
11    pub input: Vec<String>,
12}
13
14#[derive(Deserialize, Serialize)]
15pub struct EmbeddingResponse {
16    pub object: String,
17    pub model: String,
18    pub usage: EmbeddingResponseUsage,
19    pub data: Vec<Embedding>,
20}
21
22#[derive(Deserialize, Serialize)]
23pub struct Embedding {
24    pub embedding: Vec<f32>,
25    pub index: i32,
26    pub object: String,
27}
28
29#[derive(Deserialize, Serialize)]
30pub struct EmbeddingResponseUsage {
31    pub prompt_tokens: i32,
32    pub total_tokens: i32,
33}
34
35/// Creates an embedding vector for each string passed as an argument.
36pub async fn create_embeddings(
37    app_config: &ApplicationConfiguration,
38    inputs: Vec<String>,
39) -> UtilResult<Vec<Vec<f32>>> {
40    let app_config = app_config.to_owned();
41    let azure_config = app_config.azure_configuration.ok_or_else(|| {
42        util_err!(
43            EmbeddingRequestBuildError,
44            "Azure configuration is missing from the application configuration"
45        )
46    })?;
47
48    let chatbot_config = azure_config.chatbot_config.ok_or_else(|| {
49        error!("Chatbot configuration missing");
50        util_err!(
51            EmbeddingRequestBuildError,
52            "Chatbot configuration is missing from the Azure configuration"
53        )
54    })?;
55    let search_config = azure_config.search_config.ok_or_else(|| {
56        util_err!(
57            EmbeddingRequestBuildError,
58            "Azure search configuration is missing from the Azure configuration"
59        )
60    })?;
61
62    let api_endpoint = chatbot_config.embeddings_endpoint()?;
63    let input_len = inputs.len();
64    let response = REQWEST_CLIENT
65        .post(api_endpoint)
66        .header("Content-Type", "application/json")
67        .header("api-key", chatbot_config.api_key.expose_secret())
68        .json(&EmbeddingRequest {
69            model: search_config.vectorizer_model_name.to_owned(),
70            input: inputs,
71        })
72        .send()
73        .await?;
74
75    if response.status().is_success() {
76        let body = &response.text().await?;
77        let mut json: EmbeddingResponse = serde_json::from_str(body)?;
78
79        if json.data.len() != input_len {
80            return Err(util_err!(
81                EmbeddingRequestBuildError,
82                format!(
83                    "Embedding API returned {} embeddings for {} inputs",
84                    json.data.len(),
85                    input_len
86                )
87            ));
88        }
89        json.data.sort_by_key(|e| e.index);
90        let embeddings: Vec<Vec<f32>> = json.data.into_iter().map(|e| e.embedding).collect();
91
92        Ok(embeddings)
93    } else {
94        Err(util_err!(
95            EmbeddingRequestBuildError,
96            format!(
97                "Embedding API failed: {} - {}",
98                response.status(),
99                response.text().await?
100            )
101        ))
102    }
103}