Update app.py

app.py

@@ -1,8 +1,4 @@
-import os
-import sys
-import importlib
-import importlib.util
-import pkgutil
+import os, sys, importlib, importlib.util, pkgutil
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from huggingface_hub import HfApi
@@ -11,34 +7,46 @@ def load_private_code():
     api = HfApi(token=os.environ["PRIVATE_REPOSITORY_HF_TOKEN"])
 
     with TemporaryDirectory() as tmp_dir:
+        root = Path(tmp_dir)
         api.snapshot_download(
             repo_id="neoai-inc-private/Auriga",
             repo_type="space",
-            local_dir=tmp_dir,
+            local_dir=root
         )
-        root = Path(tmp_dir)
+
+        sys.path.insert(0, str(root))
+        auriga = root / "Auriga"
+        if auriga.exists():
+            sys.path.insert(0, str(auriga))
+
         app_path = (root / "app.py") if (root / "app.py").exists() else next(root.rglob("app.py"), None)
-        os.chdir(app_path.parent)
-
-        spec = importlib.util.spec_from_file_location("private_app", app_path)
-        private_module = importlib.util.module_from_spec(spec)
-        sys.modules["private_app"] = private_module
-        spec.loader.exec_module(private_module)
-
-        import neollm
-        skipped = []
-        for m in pkgutil.walk_packages(neollm.__path__, prefix="neollm."):
-            try:
-                importlib.import_module(m.name)
-            except (ModuleNotFoundError, ImportError) as e:
-                skipped.append((m.name, str(e)))
+        if not app_path:
+            raise FileNotFoundError("app.py が見つかりませんでした")
+
+        prev_cwd = os.getcwd()
         try:
+            os.chdir(app_path.parent)
+
+            spec = importlib.util.spec_from_file_location("private_app", app_path)
+            private_module = importlib.util.module_from_spec(spec)
+            sys.modules["private_app"] = private_module
+            spec.loader.exec_module(private_module)
+
+            # 事前インポート
+            import neollm
+            skipped = []
+            for m in pkgutil.walk_packages(neollm.__path__, prefix="neollm."):
+                try:
+                    importlib.import_module(m.name)
+                except (ModuleNotFoundError, ImportError) as e:
+                    skipped.append((m.name, str(e)))
+
             importlib.import_module("neollm.llm.get_llm")
             importlib.import_module("neollm.llm.model_name._abstract_model_name")
             importlib.import_module("neollm.llm.model_name.azure_model_name")
-        except Exception as e:
-            # ここが落ちる場合は、Azureに必要な依存が無い可能性があるため明示的にエラーにする
-            raise RuntimeError(f"Required preloads for Azure failed: {e}")
+
+        finally:
+            os.chdir(prev_cwd)
 
         return private_module