🐛: Add function to replace env vars when using multi-llm config

taskweaver/config/config_mgt.py

@@ -4,6 +4,7 @@
 import re
 from dataclasses import dataclass
 from typing import Any, Dict, List, Literal, NamedTuple, Optional
+from ..utils import replace_env_vars
 
 AppConfigSourceType = Literal["override", "env", "json", "app", "default"]
 AppConfigValueType = Literal["str", "int", "float", "bool", "list", "enum", "path", "dict"]
@@ -97,7 +98,10 @@ def _get_config_value(
  return val
 
  if var_name in self.json_file_store.keys():
- return self.json_file_store.get(var_name, default_value)
+ val = self.json_file_store.get(var_name, default_value)
+ # e.g., llm.api_base -> LLM_API_BASE
+ val = replace_env_vars(val)
+ return val
 
  if default_value is not None:
  return default_value

taskweaver/utils/__init__.py

@@ -11,6 +11,7 @@
 from datetime import datetime
 from hashlib import md5
 from typing import Any, Dict, List, Union
+from .json_replace import replace_env_vars
 
 
 def create_id(length: int = 4) -> str:

taskweaver/utils/json_replace.py

@@ -0,0 +1,12 @@
+import os
+
+# Function to replace env variables on nested JSON
+def replace_env_vars(data, key=None):
+ if isinstance(data, list):
+ return data
+ if isinstance(data, dict): 
+ return {k: replace_env_vars(v) for k, v in data.items()}
+ env_var_name = data.upper().replace(".", "_")
+ env_val = os.environ.get(env_var_name, None)
+ return env_val if env_val != None else data
+