headless_lms_utils/
azure_embedding.rs1use 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
35pub 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}