Update megumin_agent/bootstrap.py

megumin_agent/bootstrap.py

@@ -12,12 +12,26 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 ADK_SRC = PROJECT_ROOT / "adk-python" / "src"
 LOCAL_DATASET_DIR = PROJECT_ROOT / "data" / "processed"
 RUNTIME_DATASET_DIR = PROJECT_ROOT / "data" / "_runtime_processed"
-HF_DATASET_REPO_ID = os.getenv("MEGUMIN_HF_DATASET_REPO_ID", "Junhoee/megumin-chat")
-HF_DATASET_FILENAME = os.getenv("MEGUMIN_HF_DATASET_FILENAME", "megumin_qa_dataset.json")
+
+
+def _dataset_repo_id() -> str:
+    return os.getenv("MEGUMIN_HF_DATASET_REPO_ID", "Junhoee/megumin-chat")
+
+
+def _dataset_filename() -> str:
+    return os.getenv("MEGUMIN_HF_DATASET_FILENAME", "megumin_qa_dataset.json")
+
+
+def _index_filename() -> str:
+    return os.getenv("MEGUMIN_FAISS_INDEX_FILENAME", "megumin_questions.faiss")
+
+
+def _metadata_filename() -> str:
+    return os.getenv("MEGUMIN_FAISS_METADATA_FILENAME", "megumin_questions_meta.json")
 
 
 def bootstrap_environment() -> None:
-    load_dotenv(PROJECT_ROOT / ".env")
+    load_dotenv(PROJECT_ROOT / ".env", override=True)
     if ADK_SRC.exists():
         adk_src = str(ADK_SRC)
         if adk_src not in sys.path:
@@ -29,14 +43,26 @@ def resolve_dataset_dir() -> Path:
 
     try:
         hf_token = os.getenv("HF_TOKEN") or None
-        downloaded_path = hf_hub_download(
-            repo_id=os.getenv("MEGUMIN_HF_DATASET_REPO_ID", HF_DATASET_REPO_ID),
-            repo_type="dataset",
-            filename=os.getenv("MEGUMIN_HF_DATASET_FILENAME", HF_DATASET_FILENAME),
-            token=hf_token,
-            local_dir=str(RUNTIME_DATASET_DIR),
+        repo_id = _dataset_repo_id()
+        artifact_names = (
+            _dataset_filename(),
+            _index_filename(),
+            _metadata_filename(),
         )
-        return Path(downloaded_path).parent
+        for artifact_name in artifact_names:
+            try:
+                hf_hub_download(
+                    repo_id=repo_id,
+                    repo_type="dataset",
+                    filename=artifact_name,
+                    token=hf_token,
+                    local_dir=str(RUNTIME_DATASET_DIR),
+                )
+            except Exception:
+                if artifact_name != _dataset_filename():
+                    continue
+                raise
+        return RUNTIME_DATASET_DIR
     except Exception:
         if LOCAL_DATASET_DIR.exists() and any(LOCAL_DATASET_DIR.glob("*.json")):
             return LOCAL_DATASET_DIR