1use std::collections::HashMap;
8use std::path::PathBuf;
9
10use serde_json::Value;
11
12use super::traits::ResourceKind;
13
14#[derive(Debug, Clone)]
16pub struct ManagedResources {
17 pub knowledge_source_name: String,
18 pub index: Option<String>,
19 pub indexer: Option<String>,
20 pub datasource: Option<String>,
21 pub skillset: Option<String>,
22}
23
24pub type ManagedMap = HashMap<(ResourceKind, String), String>;
26
27pub fn extract_managed_resources(ks_name: &str, ks_def: &Value) -> ManagedResources {
32 let mut managed = ManagedResources {
33 knowledge_source_name: ks_name.to_string(),
34 index: None,
35 indexer: None,
36 datasource: None,
37 skillset: None,
38 };
39
40 let created = ks_def.get("createdResources").or_else(|| {
42 for key in &[
44 "azureBlobParameters",
45 "azureTableParameters",
46 "sharePointParameters",
47 ] {
48 if let Some(params) = ks_def.get(key) {
49 if let Some(cr) = params.get("createdResources") {
50 return Some(cr);
51 }
52 }
53 }
54 None
55 });
56
57 if let Some(created) = created {
58 if let Some(obj) = created.as_object() {
59 managed.index = obj.get("index").and_then(|v| v.as_str()).map(String::from);
60 managed.indexer = obj
61 .get("indexer")
62 .and_then(|v| v.as_str())
63 .map(String::from);
64 managed.datasource = obj
65 .get("datasource")
66 .or_else(|| obj.get("dataSource"))
67 .and_then(|v| v.as_str())
68 .map(String::from);
69 managed.skillset = obj
70 .get("skillset")
71 .and_then(|v| v.as_str())
72 .map(String::from);
73 }
74 }
75
76 managed
77}
78
79pub fn build_managed_map(knowledge_sources: &[(String, Value)]) -> ManagedMap {
84 let mut map = ManagedMap::new();
85
86 for (ks_name, ks_def) in knowledge_sources {
87 let managed = extract_managed_resources(ks_name, ks_def);
88
89 if let Some(ref name) = managed.index {
90 map.insert((ResourceKind::Index, name.clone()), ks_name.clone());
91 }
92 if let Some(ref name) = managed.indexer {
93 map.insert((ResourceKind::Indexer, name.clone()), ks_name.clone());
94 }
95 if let Some(ref name) = managed.datasource {
96 map.insert((ResourceKind::DataSource, name.clone()), ks_name.clone());
97 }
98 if let Some(ref name) = managed.skillset {
99 map.insert((ResourceKind::Skillset, name.clone()), ks_name.clone());
100 }
101 }
102
103 map
104}
105
106pub fn managing_ks<'a>(map: &'a ManagedMap, kind: ResourceKind, name: &str) -> Option<&'a String> {
108 map.get(&(kind, name.to_string()))
109}
110
111pub fn resource_directory(kind: ResourceKind, name: &str, map: &ManagedMap) -> PathBuf {
117 if kind == ResourceKind::KnowledgeSource {
118 PathBuf::from("agentic-retrieval/knowledge-sources").join(name)
120 } else if let Some(ks_name) = managing_ks(map, kind, name) {
121 PathBuf::from("agentic-retrieval/knowledge-sources").join(ks_name)
123 } else {
124 PathBuf::from(kind.directory_name())
126 }
127}
128
129pub fn resource_filename(kind: ResourceKind, name: &str, map: &ManagedMap) -> String {
135 if kind == ResourceKind::KnowledgeSource {
136 format!("{}.json", name)
138 } else if let Some(ks_name) = managing_ks(map, kind, name) {
139 let suffix = match kind {
141 ResourceKind::Index => "index",
142 ResourceKind::Indexer => "indexer",
143 ResourceKind::DataSource => "datasource",
144 ResourceKind::Skillset => "skillset",
145 _ => "resource",
146 };
147 format!("{}-{}.json", ks_name, suffix)
148 } else {
149 format!("{}.json", name)
151 }
152}
153
154pub fn find_kb_references(ks_name: &str, kbs: &[(String, Value)]) -> Vec<String> {
156 let mut refs = Vec::new();
157
158 for (kb_name, kb_def) in kbs {
159 if let Some(sources) = kb_def.get("knowledgeSources").and_then(|v| v.as_array()) {
160 for source in sources {
161 if let Some(source_name) = source
162 .as_object()
163 .and_then(|o| o.get("name"))
164 .and_then(|n| n.as_str())
165 {
166 if source_name == ks_name {
167 refs.push(kb_name.clone());
168 break;
169 }
170 }
171 }
172 }
173 }
174
175 refs
176}
177
178pub const MANAGED_SUB_RESOURCE_KINDS: &[ResourceKind] = &[
180 ResourceKind::Index,
181 ResourceKind::Indexer,
182 ResourceKind::DataSource,
183 ResourceKind::Skillset,
184];
185
186pub fn read_managed_sub_resources(
191 ks_dir: &std::path::Path,
192 ks_name: &str,
193) -> Vec<(ResourceKind, String, Value)> {
194 let mut results = Vec::new();
195
196 let suffixes = [
197 ("index", ResourceKind::Index),
198 ("indexer", ResourceKind::Indexer),
199 ("datasource", ResourceKind::DataSource),
200 ("skillset", ResourceKind::Skillset),
201 ];
202
203 for (suffix, kind) in &suffixes {
204 let filename = format!("{}-{}.json", ks_name, suffix);
205 let path = ks_dir.join(&filename);
206 if path.exists() {
207 if let Ok(content) = std::fs::read_to_string(&path) {
208 if let Ok(value) = serde_json::from_str::<Value>(&content) {
209 let azure_name = value
210 .get("name")
211 .and_then(|n| n.as_str())
212 .unwrap_or("")
213 .to_string();
214 results.push((*kind, azure_name, value));
215 }
216 }
217 }
218 }
219
220 results
221}
222
223#[cfg(test)]
224mod tests {
225 use super::*;
226 use serde_json::json;
227
228 #[test]
229 fn test_extract_managed_resources_top_level() {
230 let ks_def = json!({
231 "name": "test-ks",
232 "createdResources": {
233 "index": "test-ks-index",
234 "indexer": "test-ks-indexer",
235 "dataSource": "test-ks-datasource",
236 "skillset": "test-ks-skillset"
237 }
238 });
239
240 let managed = extract_managed_resources("test-ks", &ks_def);
241 assert_eq!(managed.knowledge_source_name, "test-ks");
242 assert_eq!(managed.index.as_deref(), Some("test-ks-index"));
243 assert_eq!(managed.indexer.as_deref(), Some("test-ks-indexer"));
244 assert_eq!(managed.datasource.as_deref(), Some("test-ks-datasource"));
245 assert_eq!(managed.skillset.as_deref(), Some("test-ks-skillset"));
246 }
247
248 #[test]
249 fn test_extract_managed_resources_nested() {
250 let ks_def = json!({
251 "name": "test-ks",
252 "azureBlobParameters": {
253 "containerName": "docs",
254 "createdResources": {
255 "index": "test-ks-index",
256 "indexer": "test-ks-indexer",
257 "dataSource": "test-ks-datasource",
258 "skillset": "test-ks-skillset"
259 }
260 }
261 });
262
263 let managed = extract_managed_resources("test-ks", &ks_def);
264 assert_eq!(managed.index.as_deref(), Some("test-ks-index"));
265 assert_eq!(managed.indexer.as_deref(), Some("test-ks-indexer"));
266 }
267
268 #[test]
269 fn test_extract_managed_resources_no_created() {
270 let ks_def = json!({
271 "name": "test-ks",
272 "indexName": "my-idx"
273 });
274
275 let managed = extract_managed_resources("test-ks", &ks_def);
276 assert!(managed.index.is_none());
277 assert!(managed.indexer.is_none());
278 assert!(managed.datasource.is_none());
279 assert!(managed.skillset.is_none());
280 }
281
282 #[test]
283 fn test_extract_managed_resources_partial() {
284 let ks_def = json!({
285 "name": "test-ks",
286 "createdResources": {
287 "index": "test-ks-index"
288 }
289 });
290
291 let managed = extract_managed_resources("test-ks", &ks_def);
292 assert_eq!(managed.index.as_deref(), Some("test-ks-index"));
293 assert!(managed.indexer.is_none());
294 }
295
296 #[test]
297 fn test_build_managed_map() {
298 let knowledge_sources = vec![
299 (
300 "ks-1".to_string(),
301 json!({
302 "name": "ks-1",
303 "createdResources": {
304 "index": "ks-1-index",
305 "indexer": "ks-1-indexer",
306 "dataSource": "ks-1-datasource",
307 "skillset": "ks-1-skillset"
308 }
309 }),
310 ),
311 (
312 "ks-2".to_string(),
313 json!({
314 "name": "ks-2",
315 "createdResources": {
316 "index": "ks-2-index"
317 }
318 }),
319 ),
320 ];
321
322 let map = build_managed_map(&knowledge_sources);
323
324 assert_eq!(
325 managing_ks(&map, ResourceKind::Index, "ks-1-index"),
326 Some(&"ks-1".to_string())
327 );
328 assert_eq!(
329 managing_ks(&map, ResourceKind::Indexer, "ks-1-indexer"),
330 Some(&"ks-1".to_string())
331 );
332 assert_eq!(
333 managing_ks(&map, ResourceKind::Index, "ks-2-index"),
334 Some(&"ks-2".to_string())
335 );
336 assert_eq!(
338 managing_ks(&map, ResourceKind::Index, "standalone-idx"),
339 None
340 );
341 }
342
343 #[test]
344 fn test_resource_directory_managed() {
345 let mut map = ManagedMap::new();
346 map.insert(
347 (ResourceKind::Index, "ks-1-index".to_string()),
348 "ks-1".to_string(),
349 );
350
351 assert_eq!(
352 resource_directory(ResourceKind::Index, "ks-1-index", &map),
353 PathBuf::from("agentic-retrieval/knowledge-sources/ks-1")
354 );
355 }
356
357 #[test]
358 fn test_resource_directory_standalone() {
359 let map = ManagedMap::new();
360 assert_eq!(
361 resource_directory(ResourceKind::Index, "my-index", &map),
362 PathBuf::from("search-management/indexes")
363 );
364 }
365
366 #[test]
367 fn test_resource_directory_ks() {
368 let map = ManagedMap::new();
369 assert_eq!(
370 resource_directory(ResourceKind::KnowledgeSource, "test-ks", &map),
371 PathBuf::from("agentic-retrieval/knowledge-sources/test-ks")
372 );
373 }
374
375 #[test]
376 fn test_resource_filename_managed() {
377 let mut map = ManagedMap::new();
378 map.insert(
379 (ResourceKind::Index, "ks-1-index".to_string()),
380 "ks-1".to_string(),
381 );
382 map.insert(
383 (ResourceKind::Indexer, "ks-1-indexer".to_string()),
384 "ks-1".to_string(),
385 );
386 map.insert(
387 (ResourceKind::DataSource, "ks-1-datasource".to_string()),
388 "ks-1".to_string(),
389 );
390 map.insert(
391 (ResourceKind::Skillset, "ks-1-skillset".to_string()),
392 "ks-1".to_string(),
393 );
394
395 assert_eq!(
396 resource_filename(ResourceKind::Index, "ks-1-index", &map),
397 "ks-1-index.json"
398 );
399 assert_eq!(
400 resource_filename(ResourceKind::Indexer, "ks-1-indexer", &map),
401 "ks-1-indexer.json"
402 );
403 assert_eq!(
404 resource_filename(ResourceKind::DataSource, "ks-1-datasource", &map),
405 "ks-1-datasource.json"
406 );
407 assert_eq!(
408 resource_filename(ResourceKind::Skillset, "ks-1-skillset", &map),
409 "ks-1-skillset.json"
410 );
411 }
412
413 #[test]
414 fn test_resource_filename_standalone() {
415 let map = ManagedMap::new();
416 assert_eq!(
417 resource_filename(ResourceKind::Index, "my-index", &map),
418 "my-index.json"
419 );
420 }
421
422 #[test]
423 fn test_resource_filename_ks() {
424 let map = ManagedMap::new();
425 assert_eq!(
426 resource_filename(ResourceKind::KnowledgeSource, "test-ks", &map),
427 "test-ks.json"
428 );
429 }
430
431 #[test]
432 fn test_find_kb_references() {
433 let kbs = vec![
434 (
435 "kb-1".to_string(),
436 json!({
437 "name": "kb-1",
438 "knowledgeSources": [
439 {"name": "ks-1"},
440 {"name": "ks-2"}
441 ]
442 }),
443 ),
444 (
445 "kb-2".to_string(),
446 json!({
447 "name": "kb-2",
448 "knowledgeSources": [
449 {"name": "ks-3"}
450 ]
451 }),
452 ),
453 ];
454
455 let refs = find_kb_references("ks-1", &kbs);
456 assert_eq!(refs, vec!["kb-1"]);
457
458 let refs = find_kb_references("ks-3", &kbs);
459 assert_eq!(refs, vec!["kb-2"]);
460
461 let refs = find_kb_references("ks-missing", &kbs);
462 assert!(refs.is_empty());
463 }
464
465 #[test]
466 fn test_find_kb_references_no_sources_array() {
467 let kbs = vec![(
468 "kb-1".to_string(),
469 json!({
470 "name": "kb-1"
471 }),
472 )];
473
474 let refs = find_kb_references("ks-1", &kbs);
475 assert!(refs.is_empty());
476 }
477
478 #[test]
479 fn test_build_managed_map_empty() {
480 let map = build_managed_map(&[]);
481 assert!(map.is_empty());
482 }
483
484 #[test]
485 fn test_read_managed_sub_resources() {
486 let dir = tempfile::tempdir().unwrap();
487 let ks_dir = dir.path().join("test-ks");
488 std::fs::create_dir_all(&ks_dir).unwrap();
489
490 std::fs::write(
492 ks_dir.join("test-ks-index.json"),
493 r#"{"name": "test-ks-index", "fields": []}"#,
494 )
495 .unwrap();
496
497 std::fs::write(
499 ks_dir.join("test-ks-skillset.json"),
500 r#"{"name": "test-ks-skillset", "skills": []}"#,
501 )
502 .unwrap();
503
504 let results = read_managed_sub_resources(&ks_dir, "test-ks");
505 assert_eq!(results.len(), 2);
506
507 let index = results.iter().find(|(k, _, _)| *k == ResourceKind::Index);
508 assert!(index.is_some());
509 assert_eq!(index.unwrap().1, "test-ks-index");
510
511 let skillset = results
512 .iter()
513 .find(|(k, _, _)| *k == ResourceKind::Skillset);
514 assert!(skillset.is_some());
515 assert_eq!(skillset.unwrap().1, "test-ks-skillset");
516 }
517
518 #[test]
519 fn test_read_managed_sub_resources_empty_dir() {
520 let dir = tempfile::tempdir().unwrap();
521 let ks_dir = dir.path().join("test-ks");
522 std::fs::create_dir_all(&ks_dir).unwrap();
523
524 let results = read_managed_sub_resources(&ks_dir, "test-ks");
525 assert!(results.is_empty());
526 }
527
528 #[test]
529 fn test_multiple_ks_managed_map() {
530 let knowledge_sources = vec![
531 (
532 "ks-a".to_string(),
533 json!({
534 "name": "ks-a",
535 "createdResources": {
536 "index": "ks-a-index",
537 "indexer": "ks-a-indexer"
538 }
539 }),
540 ),
541 (
542 "ks-b".to_string(),
543 json!({
544 "name": "ks-b",
545 "createdResources": {
546 "index": "ks-b-index",
547 "indexer": "ks-b-indexer"
548 }
549 }),
550 ),
551 ];
552
553 let map = build_managed_map(&knowledge_sources);
554
555 assert_eq!(
556 managing_ks(&map, ResourceKind::Index, "ks-a-index"),
557 Some(&"ks-a".to_string())
558 );
559 assert_eq!(
560 managing_ks(&map, ResourceKind::Index, "ks-b-index"),
561 Some(&"ks-b".to_string())
562 );
563 assert_ne!(
565 managing_ks(&map, ResourceKind::Index, "ks-a-index"),
566 Some(&"ks-b".to_string())
567 );
568 }
569}