File indexing completed on 2026-09-01 09:34:30
0001 import os
0002 import tomllib
0003
0004
0005 class PromptProcessingConfigMixin:
0006 """Shared prompt-processing config helpers for local testbed agents."""
0007
0008 def _prompt_processing_config_path(self):
0009 return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "workflows", "prompt_processing.toml")
0010
0011 def _load_prompt_processing_section(self, config_path, warn=False):
0012 if not config_path:
0013 return {}
0014 try:
0015 with open(config_path, "rb") as config_file:
0016 return tomllib.load(config_file).get("prompt_processing", {})
0017 except (OSError, TypeError, tomllib.TOMLDecodeError) as e:
0018 if warn:
0019 self.logger.warning(
0020 f"Could not load prompt_processing config from {config_path}: {e}",
0021 extra=self._log_extra()
0022 )
0023 return {}
0024
0025 def _load_prompt_processing_config(self):
0026 """Load prompt-processing settings, with workflow defaults plus active config overrides."""
0027 prompt_config = self._load_prompt_processing_section(self._prompt_processing_config_path(), warn=True)
0028 active_config = self._load_prompt_processing_section(self.config_path, warn=True)
0029 prompt_config.update(active_config)
0030 return prompt_config
0031
0032 def _config_bool(self, config, key, env_var, default):
0033 """Read a boolean setting from config, with an environment override."""
0034 value = os.getenv(env_var, config.get(key, default))
0035 if isinstance(value, bool):
0036 return value
0037 if isinstance(value, str):
0038 return value.strip().lower() in {"1", "true", "yes", "on"}
0039 return bool(value)
0040
0041 def _config_int(self, config, key, env_var, default):
0042 """Read an integer setting from config, with an environment override."""
0043 value = os.getenv(env_var, config.get(key, default))
0044 try:
0045 return int(value)
0046 except (TypeError, ValueError):
0047 self.logger.warning(
0048 f"Invalid {key} value {value!r}; using default {default}",
0049 extra=self._log_extra()
0050 )
0051 return default
0052
0053 def _config_list(self, config, key, env_var, default):
0054 """Read a comma-separated list setting from config, with an environment override."""
0055 value = os.getenv(env_var, config.get(key, default))
0056 if isinstance(value, str):
0057 return [item.strip() for item in value.split(",") if item.strip()]
0058 if isinstance(value, (list, tuple)):
0059 return [str(item).strip() for item in value if str(item).strip()]
0060 return list(default)
0061
0062 def _message_bool(self, message_data, key, default):
0063 value = message_data.get(key, default)
0064 if isinstance(value, bool):
0065 return value
0066 if isinstance(value, str):
0067 return value.strip().lower() in {"1", "true", "yes", "on"}
0068 return bool(value)
0069
0070
0071 class DecisionDatasetNamingMixin:
0072 """Shared decision-box message and dataset helpers."""
0073
0074 def _decision_box_context_for_run(self, run_id):
0075 return {}
0076
0077 def _decision_box_enabled_for_message(self, message_data, run_id=None):
0078 if "decision_box_enabled" in message_data:
0079 return self._message_bool(message_data, "decision_box_enabled", self.decision_box_enabled)
0080 if run_id is not None:
0081 context = self._decision_box_context_for_run(run_id)
0082 if "decision_box_enabled" in context:
0083 return bool(context["decision_box_enabled"])
0084 return self.decision_box_enabled
0085
0086 def _non_decision_box_site_for_message(self, message_data, run_id=None):
0087 site = message_data.get("non_decision_box_site")
0088 if site:
0089 return str(site).strip()
0090 if run_id is not None:
0091 context = self._decision_box_context_for_run(run_id)
0092 site = context.get("non_decision_box_site")
0093 if site:
0094 return str(site).strip()
0095 return getattr(self, "non_decision_box_site", None)
0096
0097 def _run_dataset_name(self, run_number=None):
0098 dataset = getattr(self, "dataset", None)
0099 if dataset:
0100 return dataset
0101 if run_number is not None:
0102 return f"swf.{run_number}.run"
0103 return ""
0104
0105 def _run_dataset_did(self, run_number=None):
0106 return f"{self.decision_box_rucio_scope}:{self._run_dataset_name(run_number)}"
0107
0108 def _input_dataset_name_for_site(self, run_number, site_name):
0109 run_dataset_name = f"swf.{run_number}.run"
0110 if self.decision_box_site_dataset_template:
0111 return self.decision_box_site_dataset_template.format(
0112 run_dataset_name=run_dataset_name,
0113 run_number=run_number,
0114 site_name=site_name,
0115 site=site_name,
0116 )
0117 return f"{self._run_dataset_name(run_number)}.{site_name}"
0118
0119 def _input_dataset_did_for_site(self, run_number, site_name):
0120 return f"{self.decision_box_rucio_scope}:{self._input_dataset_name_for_site(run_number, site_name)}"
0121
0122 def _site_name_for_dataset(self, run_number, dataset_did):
0123 for site_name in self.decision_box_sites:
0124 if dataset_did == self._input_dataset_did_for_site(run_number, site_name):
0125 return site_name
0126 return None
0127
