move env-var reading prior to the model parse

microsoft · Apr 2, 2024 · e0ba225 · e0ba225
1 parent 56e582b
commit e0ba225
Show file tree

Hide file tree

Showing 10 changed files with 49 additions and 13 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -79,7 +79,6 @@ jobs:
         name: Unit Tests
         env:
           GRAPHRAG_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-          GRAPHRAG_BLOB_CONNECTION_STRING: ${{ secrets.GRAPHRAG_BLOB_CONNECTION_STRING}}
 
       - run: yarn test:integration
         name: Integration Tests
@@ -92,3 +91,4 @@ jobs:
           GRAPHRAG_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           GRAPHRAG_LLM_MODEL: gpt-3.5-turbo
           GRAPHRAG_EMBEDDING_MODEL: text-embedding-3-small
+          GRAPHRAG_BLOB_CONNECTION_STRING: ${{ secrets.GRAPHRAG_BLOB_CONNECTION_STRING}}
diff --git a/python/graphrag/graphrag/index/cli.py b/python/graphrag/graphrag/index/cli.py
@@ -12,9 +12,8 @@
 from graphrag.index.cache import NoopPipelineCache
 from graphrag.index.config import PipelineConfig
 from graphrag.index.default_config import (
-    DefaultConfigParametersModel,
     default_config,
-    default_config_parameters,
+    default_config_parameters_from_dict,
     default_config_parameters_from_env_vars,
 )
 from graphrag.index.progress import (
@@ -170,17 +169,15 @@ def _read_config_parameters(root: str, reporter: ProgressReporter):
             import yaml
 
             data = yaml.safe_load(file)
-            model = DefaultConfigParametersModel.model_validate(data)
-            return default_config_parameters(model, root)
+            return default_config_parameters_from_dict(data, root)
 
     if settings_json.exists():
         reporter.success(f"Reading settings from {settings_json}")
         with settings_json.open("r") as file:
             import json
 
             data = json.loads(file.read())
-            model = DefaultConfigParametersModel.model_validate(data)
-            return default_config_parameters(data, root)
+            return default_config_parameters_from_dict(data, root)
 
     reporter.success("Reading settings from environment variables")
     return default_config_parameters_from_env_vars(root)

diff --git a/python/graphrag/graphrag/index/default_config/__init__.py b/python/graphrag/graphrag/index/default_config/__init__.py
@@ -24,6 +24,7 @@
     TextEmbeddingConfigModel,
     UmapConfigModel,
     default_config_parameters,
+    default_config_parameters_from_dict,
     default_config_parameters_from_env_vars,
 )
 
@@ -48,6 +49,7 @@
     "UmapConfigModel",
     "default_config",
     "default_config_parameters",
+    "default_config_parameters_from_dict",
     "default_config_parameters_from_env_vars",
     "load_pipeline_config",
 ]
diff --git a/python/graphrag/graphrag/index/default_config/parameters/__init__.py b/python/graphrag/graphrag/index/default_config/parameters/__init__.py
@@ -5,6 +5,7 @@
 from .default_config_parameters_model import DefaultConfigParametersModel
 from .factories import (
     default_config_parameters,
+    default_config_parameters_from_dict,
     default_config_parameters_from_env_vars,
 )
 from .models import (
@@ -48,6 +49,7 @@
     "TextEmbeddingConfigModel",
     "UmapConfigModel",
     "default_config_parameters",
+    "default_config_parameters_from_dict",
     "default_config_parameters_from_env_vars",
     "read_dotenv",
 ]
diff --git a/python/graphrag/graphrag/index/default_config/parameters/factories.py b/python/graphrag/graphrag/index/default_config/parameters/factories.py
@@ -41,12 +41,41 @@
 )
 
 
+def perform_replacements(values: dict, env: Env):
+    """Perform env-var replacements in the dictionary."""
+
+    def traverse_dict(dictionary):
+        for key, value in dictionary.items():
+            if isinstance(value, dict):
+                traverse_dict(value)
+            elif (
+                isinstance(value, str)
+                and value.startswith("${")
+                and value.endswith("}")
+            ):
+                env_var = value[2:-1]
+                dictionary[key] = env(env_var)
+
+    traverse_dict(values)
+
+
+def default_config_parameters_from_dict(values: dict, root_dir: str | None):
+    """Load Configuration Parameters from a dictionary."""
+    root_dir = root_dir or str(Path.cwd())
+    env = _make_env(root_dir)
+    perform_replacements(values, env)
+
+    model = DefaultConfigParametersModel.model_validate(values)
+    return default_config_parameters(model, root_dir)
+
+
 def default_config_parameters(
     values: DefaultConfigParametersModel, root_dir: str | None
 ):
     """Load Configuration Parameters from a dictionary."""
     root_dir = root_dir or str(Path.cwd())
     env = _make_env(root_dir)
+
     return DefaultConfigParametersDict(values, env, root_dir)
 
 

diff --git a/python/graphrag/graphrag/index/default_config/parameters/sections/cache_config_section.py b/python/graphrag/graphrag/index/default_config/parameters/sections/cache_config_section.py
@@ -29,7 +29,8 @@ def __init__(self, values: CacheConfigModel, env: Env):
     @property
     def type(self) -> PipelineCacheType:
         """The cache type to use."""
-        return self.replace(self._values.type, DEFAULT_CACHE_TYPE)
+        result = self.replace(str(self._values.type))
+        return PipelineCacheType(result) if result else DEFAULT_CACHE_TYPE
 
     @property
     def connection_string(self) -> str:

diff --git a/python/graphrag/graphrag/index/default_config/parameters/sections/input_config_section.py b/python/graphrag/graphrag/index/default_config/parameters/sections/input_config_section.py
@@ -34,7 +34,8 @@ def __init__(self, values: InputConfigModel, env: Env):
     @property
     def type(self) -> PipelineInputType:
         """The input type to use."""
-        return self.replace(self._values.type, DEFAULT_INPUT_TYPE)
+        result = self.replace(str(self._values.type))
+        return PipelineInputType(result) if result else DEFAULT_INPUT_TYPE
 
     @property
     def storage_type(self) -> PipelineInputStorageType:

diff --git a/...on/graphrag/graphrag/index/default_config/parameters/sections/reporting_config_section.py b/...on/graphrag/graphrag/index/default_config/parameters/sections/reporting_config_section.py
@@ -29,7 +29,8 @@ def __init__(self, values: ReportingConfigModel, env: Env):
     @property
     def type(self) -> PipelineReportingType:
         """The reporting type to use."""
-        return self.replace(self._values.type, DEFAULT_REPORTING_TYPE)
+        result = self.replace(str(self._values.type))
+        return PipelineReportingType(result) if result else DEFAULT_REPORTING_TYPE
 
     @property
     def connection_string(self) -> str:

diff --git a/python/graphrag/graphrag/index/default_config/parameters/sections/storage_config_section.py b/python/graphrag/graphrag/index/default_config/parameters/sections/storage_config_section.py
@@ -29,7 +29,8 @@ def __init__(self, values: StorageConfigModel, env: Env):
     @property
     def type(self) -> PipelineStorageType:
         """The storage type to use."""
-        return self.replace(self._values.type, DEFAULT_STORAGE_TYPE)
+        result = self.replace(str(self._values.type))
+        return PipelineStorageType(result) if result else DEFAULT_STORAGE_TYPE
 
     @property
     def connection_string(self) -> str:

diff --git a/python/graphrag/tests/smoke/test_fixtures.py b/python/graphrag/tests/smoke/test_fixtures.py
@@ -24,6 +24,8 @@
 BLOB_CONNECTION_STRING = os.environ.get(
     "GRAPHRAG_BLOB_CONNECTION_STRING", WELL_KNOWN_AZURITE_CONNECTION_STRING
 )
+CACHE_TYPE = os.environ.get("CACHE_TYPE", "file") or "file"
+CACHE_CONTAINER_NAME = os.environ.get("CACHE_CONTAINER_NAME", "cache") or "cache"
 
 
 def _load_fixtures():
@@ -244,8 +246,8 @@ def __run_query(self, root: Path, query_config: dict[str, str]):
         os.environ,
         {
             **os.environ,
-            "CACHE_TYPE": os.environ.get("CACHE_TYPE", "file"),
-            "CACHE_CONTAINER_NAME": os.environ.get("CACHE_CONTAINER_NAME", "cache"),
+            "CACHE_TYPE": CACHE_TYPE,
+            "CACHE_CONTAINER_NAME": CACHE_CONTAINER_NAME,
             "BLOB_STORAGE_CONNECTION_STRING": BLOB_CONNECTION_STRING,
         },
         clear=True,