Update server.py

server.py

@@ -20,18 +20,21 @@ from graphrag.query.structured_search.local_search.mixed_context import (
 from graphrag.query.structured_search.local_search.search import LocalSearch
 from graphrag.vector_stores.lancedb import LanceDBVectorStore
 
+# 基础数据目录路径
+BASE_DATA_DIR = "/app/graphrag-data/data"
 
-# 定义不同数据集的配置
-DATA_CONFIGS = {
-    "ghost": {
-        "input_dir": "/app/graphrag-data/data/the_bit_player",
-        "community_level": 2
-    },
-    "zhu_rongji": {
-        "input_dir": "/app/graphrag-data/data/the_bit_player", 
-        "community_level": 2
+# 初始化DATA_CONFIGS字典
+DATA_CONFIGS = {}
+
+# 获取目录下所有文件夹
+data_dirs = [d for d in os.listdir(BASE_DATA_DIR) if os.path.isdir(os.path.join(BASE_DATA_DIR, d))]
+
+# 为每个文件夹创建配置
+for dir_name in data_dirs:
+    DATA_CONFIGS[dir_name] = {
+        "input_dir": os.path.join(BASE_DATA_DIR, dir_name),
+        "community_level": 2  # 默认值设置为2
     }
-}
 
 api_key = os.environ['api_key']
 llm_model = os.environ['llm_model']