Update loader.py (#32)

fix import bugs
Gentopia-AI · Jul 9, 2023 · 36a7139 · 36a7139
1 parent e51eec6
commit 36a7139
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/gentopia/assembler/loader.py b/gentopia/assembler/loader.py
@@ -1,3 +1,4 @@
+import importlib
 from pathlib import Path
 from typing import Any, IO
 from gentopia.prompt import *
@@ -26,7 +27,12 @@ def include(self, node: yaml.Node) -> Any:
 
     def prompt(self, node: yaml.Node) -> Any:
         prompt = self.construct_scalar(node)
-        prompt_cls = eval(prompt)
+        if '.' in prompt:
+            _path = prompt.split('.')
+            module = importlib.import_module('.'.join(_path[:-1]))
+            prompt_cls = getattr(module, _path[-1])
+        else:
+            prompt_cls = eval(prompt)
         assert issubclass(prompt_cls.__class__, PromptTemplate)
         return prompt_cls