Coverage for agentlib/utils/plugin_import.py: 54%

1""" 

2Module containing all function to import new plugins 

3""" 

4 

5import importlib 

6from pydantic import BaseModel 

7 

8 

9class ModuleImport(BaseModel): 

10 """ 

11 Data-Class to import a given python file 

12 from ``import_path`` and load the given 

13 ``class_name`` 

14 """ 

15 

16 import_path: str 

17 class_name: str 

18 

19 def import_class(self): 

20 """Import the Module with class_name from the import path""" 

21 module = importlib.import_module(self.import_path) 

22 return getattr(module, self.class_name) 

23 

24 

25class SaveUpdateDict(dict): 

26 """ 

27 Custom object to safely update the dictionary. 

28 Duplicate entries will raise an error. 

29 """ 

30 

31 def update(self, new_dict, **kwargs) -> None: 

32 """Check if all modules have distinct identifier strings 

33 

34 Args: 

35 new_dict: 

36 **kwargs: 

37 """ 

38 duplicates = set(self.keys()).intersection(new_dict.keys()) 

39 if duplicates: 

40 raise KeyError( 

41 "One or multiple MODULE_TYPES contain " 

42 "the same string identifier (type). " 

43 f"The type has to be unique! {', '.join(duplicates)}" 

44 ) 

45 super().update(new_dict) 

46 

47 

48def load_plugin(name: str, loaded_classes: SaveUpdateDict, plugin_types_name: str): 

49 """ 

50 Loads the plugin based on the given name. 

51 

52 Args: 

53 name str: 

54 Name of the plugin 

55 loaded_classes SaveUpdateDict: 

56 SaveUpdateDict instance with already 

57 loaded classes (modules or models) 

58 plugin_types_name str: 

59 Name of the dictionary in the plugin. 

60 Typical values are "MODULE_TYPES" or "MODEL_TYPES". 

61 """ 

62 for key in loaded_classes: 

63 if key.startswith(f"{name}."): 

64 return # Already loaded 

65 try: 

66 plugin = importlib.import_module(name) 

67 except ImportError as err: 

68 raise ImportError(f"Plugin '{name}' is not installed.") from err 

69 try: 

70 plugin_dict = getattr(plugin, plugin_types_name) 

71 except AttributeError: 

72 raise ImportError( 

73 f"Plugin '{name}' has no dictionary called " 

74 f"'{plugin_types_name}' to import plugin types." 

75 ) 

76 if not isinstance(plugin_dict, dict): 

77 raise TypeError( 

78 f"Loaded object '{plugin_types_name}' is not a dictionary " 

79 f"but a {type(plugin_dict)}" 

80 ) 

81 for key, value in plugin_dict.items(): 

82 loaded_classes.update({f"{name}.{key}": value})