1use secrecy::ExposeSecret;
2
3use crate::{
4 llm_utils::{azure_search_configuration, azure_search_request},
5 prelude::*,
6};
7use headless_lms_base::config::ApplicationConfiguration;
8
9const API_VERSION: &str = "2024-07-01";
10
11#[derive(Serialize, Deserialize)]
12#[serde(rename_all = "camelCase")]
13pub struct NewIndex {
14 name: String,
15 fields: Vec<Field>,
16 scoring_profiles: Vec<ScoringProfile>,
17 default_scoring_profile: Option<String>,
18 suggesters: Vec<Suggester>,
19 analyzers: Vec<Analyzer>,
20 tokenizers: Vec<serde_json::Value>,
21 token_filters: Vec<serde_json::Value>,
22 char_filters: Vec<serde_json::Value>,
23 cors_options: CorsOptions,
24 encryption_key: Option<EncryptionKey>,
25 similarity: Similarity,
26 semantic: Semantic,
27 vector_search: VectorSearch,
28}
29
30#[derive(Serialize, Deserialize)]
31#[serde(rename_all = "camelCase")]
32pub struct Analyzer {
33 name: String,
34 #[serde(rename = "@odata.type")]
35 odata_type: String,
36 char_filters: Vec<String>,
37 tokenizer: String,
38}
39
40#[derive(Serialize, Deserialize)]
41#[serde(rename_all = "camelCase")]
42pub struct CorsOptions {
43 allowed_origins: Vec<String>,
44 max_age_in_seconds: i64,
45}
46
47#[derive(Serialize, Deserialize)]
48#[serde(rename_all = "camelCase")]
49pub struct EncryptionKey {
50 key_vault_key_name: String,
51 key_vault_key_version: String,
52 key_vault_uri: String,
53 access_credentials: AccessCredentials,
54}
55
56#[derive(Serialize, Deserialize)]
57#[serde(rename_all = "camelCase")]
58pub struct AccessCredentials {
59 application_id: String,
60 application_secret: String,
61}
62
63#[derive(Serialize, Deserialize)]
64#[serde(rename_all = "camelCase")]
65pub struct Field {
66 name: String,
67 #[serde(rename = "type")]
68 field_type: String,
69 key: Option<bool>,
70 searchable: Option<bool>,
71 filterable: Option<bool>,
72 sortable: Option<bool>,
73 facetable: Option<bool>,
74 retrievable: Option<bool>,
75 index_analyzer: Option<String>,
76 search_analyzer: Option<String>,
77 analyzer: Option<String>,
78 synonym_maps: Option<Vec<String>>,
79 dimensions: Option<i64>,
80 vector_search_profile: Option<String>,
81 stored: Option<bool>,
82 vector_encoding: Option<serde_json::Value>,
83}
84
85#[derive(Serialize, Deserialize)]
86pub struct ScoringProfile {
87 name: String,
88 text: Text,
89 functions: Vec<Function>,
90}
91
92#[derive(Serialize, Deserialize)]
93#[serde(rename_all = "camelCase")]
94pub struct Function {
95 #[serde(rename = "type")]
96 function_type: String,
97 boost: i64,
98 field_name: String,
99 interpolation: String,
100 distance: Distance,
101}
102
103#[derive(Serialize, Deserialize)]
104#[serde(rename_all = "camelCase")]
105pub struct Distance {
106 reference_point_parameter: String,
107 boosting_distance: i64,
108}
109
110#[derive(Serialize, Deserialize)]
111pub struct Text {
112 weights: Weights,
113}
114
115#[derive(Serialize, Deserialize)]
116#[serde(rename_all = "camelCase")]
117pub struct Weights {
118 hotel_name: i64,
119}
120
121#[derive(Serialize, Deserialize)]
122#[serde(rename_all = "camelCase")]
123pub struct Semantic {
124 default_configuration: String,
125 configurations: Vec<SemanticConfiguration>,
126}
127
128#[derive(Serialize, Deserialize)]
129#[serde(rename_all = "camelCase")]
130pub struct SemanticConfiguration {
131 name: String,
132 prioritized_fields: SemanticConfigurationPrioritizedFields,
133}
134
135#[derive(Serialize, Deserialize)]
136#[serde(rename_all = "camelCase")]
137pub struct SemanticConfigurationPrioritizedFields {
138 title_field: FieldDescriptor,
139 prioritized_content_fields: Vec<FieldDescriptor>,
140 prioritized_keywords_fields: Vec<FieldDescriptor>,
141}
142
143#[derive(Serialize, Deserialize)]
144#[serde(rename_all = "camelCase")]
145pub struct FieldDescriptor {
146 field_name: String,
147}
148
149#[derive(Serialize, Deserialize)]
150pub struct Similarity {
151 #[serde(rename = "@odata.type")]
152 odata_type: String,
153 b: Option<f64>,
154 k1: Option<f64>,
155}
156
157#[derive(Serialize, Deserialize)]
158#[serde(rename_all = "camelCase")]
159pub struct Suggester {
160 name: String,
161 search_mode: String,
162 source_fields: Vec<String>,
163}
164
165#[derive(Serialize, Deserialize)]
166pub struct VectorSearch {
167 profiles: Vec<Profile>,
168 algorithms: Vec<Algorithm>,
169 compressions: Vec<Compression>,
170 vectorizers: Vec<Vectorizer>,
171}
172
173#[derive(Serialize, Deserialize)]
174#[serde(rename_all = "camelCase")]
175pub struct Vectorizer {
176 name: String,
177 kind: String,
178 #[serde(rename = "azureOpenAIParameters")]
179 azure_open_ai_parameters: AzureOpenAiParameters,
180 custom_web_api_parameters: Option<serde_json::Value>,
181}
182
183#[derive(Serialize, Deserialize)]
184#[serde(rename_all = "camelCase")]
185pub struct AzureOpenAiParameters {
186 resource_uri: String,
187 deployment_id: String,
188 api_key: String,
189 model_name: String,
190 auth_identity: Option<serde_json::Value>,
191}
192
193#[derive(Serialize, Deserialize)]
194#[serde(rename_all = "camelCase")]
195pub struct Algorithm {
196 name: String,
197 kind: String,
198 hnsw_parameters: Option<HnswParameters>,
199 exhaustive_knn_parameters: Option<ExhaustiveKnnParameters>,
200}
201
202#[derive(Serialize, Deserialize)]
203pub struct ExhaustiveKnnParameters {
204 metric: String,
205}
206
207#[derive(Serialize, Deserialize)]
208#[serde(rename_all = "camelCase")]
209pub struct HnswParameters {
210 m: i64,
211 metric: String,
212 ef_construction: i64,
213 ef_search: i64,
214}
215
216#[derive(Serialize, Deserialize)]
217#[serde(rename_all = "camelCase")]
218pub struct Compression {
219 name: String,
220 kind: String,
221 scalar_quantization_parameters: Option<ScalarQuantizationParameters>,
222 rerank_with_original_vectors: bool,
223 default_oversampling: i64,
224}
225
226#[derive(Serialize, Deserialize)]
227#[serde(rename_all = "camelCase")]
228pub struct ScalarQuantizationParameters {
229 quantized_data_type: String,
230}
231
232#[derive(Serialize, Deserialize)]
233pub struct Profile {
234 name: String,
235 algorithm: String,
236 compression: Option<String>,
237 vectorizer: Option<String>,
238}
239
240pub async fn does_search_index_exist(
241 index_name: &str,
242 app_config: &ApplicationConfiguration,
243) -> ChatbotResult<bool> {
244 let search_config = azure_search_configuration(app_config)?;
245
246 let mut url = search_config.search_endpoint.clone();
247 url.set_path(&format!("indexes('{}')", index_name));
248 url.set_query(Some(&format!("api-version={}", API_VERSION)));
249
250 let response = azure_search_request(reqwest::Method::GET, url, search_config)
251 .send()
252 .await?;
253
254 if response.status().is_success() {
255 Ok(true)
256 } else if response.status() == 404 {
257 Ok(false)
258 } else {
259 let status = response.status();
260 let error_text = response.text().await?;
261 Err(chatbot_err!(
262 FailedAzureResponse,
263 format!(
264 "Error checking if index exists. Status: {}. Error: {}",
265 status, error_text
266 )
267 ))
268 }
269}
270
271pub async fn create_search_index(
272 index_name: String,
273 app_config: &ApplicationConfiguration,
274) -> ChatbotResult<()> {
275 let search_config = azure_search_configuration(app_config)?;
276
277 let fields = vec![
278 Field {
279 name: "chunk_id".to_string(),
280 field_type: "Edm.String".to_string(),
281 key: Some(true),
282 searchable: Some(true),
283 filterable: Some(true),
284 retrievable: Some(true),
285 stored: Some(true),
286 sortable: Some(false),
287 facetable: Some(false),
288 analyzer: Some("keyword".to_string()),
289 index_analyzer: None,
290 search_analyzer: None,
291 synonym_maps: Some(vec![]),
292 dimensions: None,
293 vector_search_profile: None,
294 vector_encoding: None,
295 },
296 Field {
297 name: "language".to_string(),
298 field_type: "Edm.String".to_string(),
299 key: Some(false),
300 searchable: Some(true),
301 filterable: Some(true),
302 retrievable: Some(true),
303 stored: Some(true),
304 sortable: Some(false),
305 facetable: Some(false),
306 analyzer: None,
307 index_analyzer: None,
308 search_analyzer: None,
309 synonym_maps: Some(vec![]),
310 dimensions: None,
311 vector_search_profile: None,
312 vector_encoding: None,
313 },
314 Field {
315 name: "parent_id".to_string(),
316 field_type: "Edm.String".to_string(),
317 key: Some(false),
318 searchable: Some(true),
319 filterable: Some(true),
320 retrievable: Some(true),
321 stored: Some(true),
322 sortable: Some(true),
323 facetable: Some(true),
324 analyzer: None,
325 index_analyzer: None,
326 search_analyzer: None,
327 synonym_maps: Some(vec![]),
328 dimensions: None,
329 vector_search_profile: None,
330 vector_encoding: None,
331 },
332 Field {
333 name: "chunk".to_string(),
334 field_type: "Edm.String".to_string(),
335 key: Some(false),
336 searchable: Some(true),
337 filterable: Some(false),
338 retrievable: Some(true),
339 stored: Some(true),
340 sortable: Some(false),
341 facetable: Some(false),
342 analyzer: Some("keyword".to_string()),
343 index_analyzer: None,
344 search_analyzer: None,
345 synonym_maps: Some(vec![]),
346 dimensions: None,
347 vector_search_profile: None,
348 vector_encoding: None,
349 },
350 Field {
351 name: "title".to_string(),
352 field_type: "Edm.String".to_string(),
353 key: Some(false),
354 searchable: Some(true),
355 filterable: Some(true),
356 retrievable: Some(true),
357 stored: Some(true),
358 sortable: Some(false),
359 facetable: Some(false),
360 analyzer: None,
361 index_analyzer: None,
362 search_analyzer: None,
363 synonym_maps: Some(vec![]),
364 dimensions: None,
365 vector_search_profile: None,
366 vector_encoding: None,
367 },
368 Field {
369 name: "url".to_string(),
370 field_type: "Edm.String".to_string(),
371 key: Some(false),
372 searchable: Some(false),
373 filterable: Some(true),
374 retrievable: Some(true),
375 stored: Some(true),
376 sortable: Some(false),
377 facetable: Some(false),
378 analyzer: None,
379 index_analyzer: None,
380 search_analyzer: None,
381 synonym_maps: Some(vec![]),
382 dimensions: None,
383 vector_search_profile: None,
384 vector_encoding: None,
385 },
386 Field {
387 name: "course_id".to_string(),
388 field_type: "Edm.String".to_string(),
389 key: Some(false),
390 searchable: Some(false),
391 filterable: Some(true),
392 retrievable: Some(true),
393 stored: Some(true),
394 sortable: Some(false),
395 facetable: Some(false),
396 analyzer: None,
397 index_analyzer: None,
398 search_analyzer: None,
399 synonym_maps: Some(vec![]),
400 dimensions: None,
401 vector_search_profile: None,
402 vector_encoding: None,
403 },
404 Field {
405 name: "text_vector".to_string(),
406 field_type: "Collection(Edm.Single)".to_string(),
407 key: Some(false),
408 searchable: Some(true),
409 filterable: Some(false),
410 retrievable: Some(true),
411 stored: Some(true),
412 sortable: Some(false),
413 facetable: Some(false),
414 analyzer: None,
415 index_analyzer: None,
416 search_analyzer: None,
417 synonym_maps: Some(vec![]),
418 dimensions: Some(1536),
419 vector_search_profile: Some(format!("{}-azureOpenAi-text-profile", index_name)),
420 vector_encoding: None,
421 },
422 Field {
423 name: "filepath".to_string(),
424 field_type: "Edm.String".to_string(),
425 key: Some(false),
426 searchable: Some(true),
427 filterable: Some(true),
428 retrievable: Some(true),
429 stored: Some(true),
430 sortable: Some(false),
431 facetable: Some(false),
432 analyzer: None,
433 index_analyzer: None,
434 search_analyzer: None,
435 synonym_maps: Some(vec![]),
436 dimensions: None,
437 vector_search_profile: None,
438 vector_encoding: None,
439 },
440 Field {
441 name: "chunk_context".to_string(),
442 field_type: "Edm.String".to_string(),
443 key: Some(false),
444 searchable: Some(true),
445 filterable: Some(false),
446 retrievable: Some(true),
447 stored: Some(true),
448 sortable: Some(false),
449 facetable: Some(false),
450 analyzer: None,
451 index_analyzer: None,
452 search_analyzer: None,
453 synonym_maps: Some(vec![]),
454 dimensions: None,
455 vector_search_profile: None,
456 vector_encoding: None,
457 },
458 ];
459
460 let index = NewIndex {
461 name: index_name.clone(),
462 fields,
463 scoring_profiles: vec![],
464 default_scoring_profile: None,
465 suggesters: vec![],
466 analyzers: vec![],
467 tokenizers: vec![],
468 token_filters: vec![],
469 char_filters: vec![],
470 cors_options: CorsOptions {
471 allowed_origins: vec!["*".to_string()],
472 max_age_in_seconds: 300,
473 },
474 encryption_key: None,
475 similarity: Similarity {
476 odata_type: "#Microsoft.Azure.Search.BM25Similarity".to_string(),
477 b: None,
478 k1: None,
479 },
480 semantic: Semantic {
481 default_configuration: format!("{}-semantic-configuration", index_name),
482 configurations: vec![SemanticConfiguration {
483 name: format!("{}-semantic-configuration", index_name),
484 prioritized_fields: SemanticConfigurationPrioritizedFields {
485 title_field: FieldDescriptor {
486 field_name: "title".to_string(),
487 },
488 prioritized_content_fields: vec![FieldDescriptor {
489 field_name: "chunk".to_string(),
490 }],
491 prioritized_keywords_fields: vec![],
492 },
493 }],
494 },
495 vector_search: VectorSearch {
496 profiles: vec![Profile {
497 name: format!("{}-azureOpenAi-text-profile", index_name),
498 algorithm: format!("{}-algorithm", index_name),
499 vectorizer: Some(format!("{}-azureOpenAi-text-vectorizer", index_name)),
500 compression: None,
501 }],
502 algorithms: vec![Algorithm {
503 name: format!("{}-algorithm", index_name),
504 kind: "hnsw".to_string(),
505 hnsw_parameters: Some(HnswParameters {
506 m: 4,
507 metric: "cosine".to_string(),
508 ef_construction: 400,
509 ef_search: 500,
510 }),
511 exhaustive_knn_parameters: None,
512 }],
513 vectorizers: vec![Vectorizer {
514 name: format!("{}-azureOpenAi-text-vectorizer", index_name),
515 kind: "azureOpenAI".to_string(),
516 azure_open_ai_parameters: AzureOpenAiParameters {
517 resource_uri: search_config.vectorizer_resource_uri.clone(),
518 deployment_id: search_config.vectorizer_deployment_id.clone(),
519 api_key: search_config.vectorizer_api_key.expose_secret().to_string(),
520 model_name: search_config.vectorizer_model_name.clone(),
521 auth_identity: None,
522 },
523 custom_web_api_parameters: None,
524 }],
525 compressions: vec![],
526 },
527 };
528
529 let index_json = serde_json::to_string(&index)?;
530
531 let mut url = search_config.search_endpoint.clone();
532 url.set_path("/indexes");
533 url.set_query(Some(&format!("api-version={}", API_VERSION)));
534
535 let response = azure_search_request(reqwest::Method::POST, url, search_config)
536 .body(index_json)
537 .send()
538 .await?;
539
540 if response.status().is_success() {
542 println!("Index created successfully: {}", index_name);
543 Ok(())
544 } else {
545 let status = response.status();
546 let error_text = response.text().await?;
547 Err(chatbot_err!(
548 FailedAzureResponse,
549 format!(
550 "Failed to create index. Status: {}. Error: {}",
551 status, error_text
552 )
553 ))
554 }
555}
556
557#[derive(Serialize, Deserialize)]
558pub struct IndexAction<T> {
559 #[serde(rename = "@search.action")]
560 pub search_action: String,
561 pub document: T,
562}
563
564#[derive(Serialize, Deserialize)]
565pub struct IndexBatch<T> {
566 pub value: Vec<IndexAction<T>>,
567}
568
569pub async fn add_documents_to_index<T>(
570 index_name: &str,
571 documents: Vec<T>,
572 app_config: &ApplicationConfiguration,
573) -> ChatbotResult<()>
574where
575 T: Serialize,
576{
577 let search_config = azure_search_configuration(app_config)?;
578
579 let mut url = search_config.search_endpoint.clone();
580 url.set_path(&format!("indexes('{}')/docs/index", index_name));
581 url.set_query(Some(&format!("api-version={}", API_VERSION)));
582
583 let index_actions: ChatbotResult<Vec<IndexAction<String>>> = documents
584 .into_iter()
585 .map(|doc| {
586 serde_json::to_string(&doc)
587 .map(|document| IndexAction {
588 search_action: "upload".to_string(),
589 document,
590 })
591 .map_err(|e| chatbot_err!(SerdeJson, "Failed to serialize document", e))
592 })
593 .collect();
594 let index_actions = index_actions?;
595
596 let batch = IndexBatch {
597 value: index_actions,
598 };
599
600 let batch_json = serde_json::to_string(&batch)?;
601
602 let response = azure_search_request(reqwest::Method::POST, url, search_config)
603 .body(batch_json)
604 .send()
605 .await?;
606
607 if response.status().is_success() {
608 println!("Documents added successfully to index: {}", index_name);
609 Ok(())
610 } else {
611 let status = response.status();
612 let error_text = response.text().await?;
613 Err(chatbot_err!(
614 FailedAzureResponse,
615 format!(
616 "Failed to add documents to index. Status: {}. Error: {}",
617 status, error_text
618 )
619 ))
620 }
621}