diff --git a/align_system/algorithms/comparative_regression_adm_component.py b/align_system/algorithms/comparative_regression_adm_component.py index 76bab085..bf0f812e 100644 --- a/align_system/algorithms/comparative_regression_adm_component.py +++ b/align_system/algorithms/comparative_regression_adm_component.py @@ -46,7 +46,8 @@ def __init__(self, target_attribute_names_override=None, enable_caching=False, reverse_choice_ordering=False, - output_conflict_resolver=None): + output_conflict_resolver=None, + choice_schema_transform=None): self.structured_inference_engine = structured_inference_engine self.scenario_description_template = scenario_description_template self.prompt_template = prompt_template @@ -69,6 +70,8 @@ def __init__(self, self.output_conflict_resolver = output_conflict_resolver + self.choice_schema_transform = choice_schema_transform + def run_returns(self): return ('attribute_prediction_reasonings', 'attribute_prediction_scores', @@ -129,6 +132,10 @@ def run(self, attribute_dialogs = {} attribute_prediction_scores = {} attribute_prediction_reasonings = {} + if self.choice_schema_transform is not None: + shortened_choices = [self.choice_schema_transform(choice) for choice in choices] + else: + shortened_choices = list(choices) for attribute in target_attributes: scenario_description = call_with_coerced_args( self.scenario_description_template, @@ -166,7 +173,7 @@ def run(self, score_schema = call_with_coerced_args( self.score_schema_template, - {'choices': choices, + {'choices': shortened_choices, 'attribute': attribute.name}) dialog_prompt = self.structured_inference_engine.dialog_to_prompt(dialog) @@ -183,16 +190,16 @@ def run(self, attribute.kdma, i), extra={"markup": True}) log.info(response, extra={"highlighter": JSON_HIGHLIGHTER}) - for choice in choices: + for choice, shortened_choice in zip(choices, shortened_choices): attribute_prediction_scores.setdefault(choice, {}) attribute_prediction_scores[choice].setdefault( - attribute.kdma, []).append(response[choice]['score'] / attribute.factor) + attribute.kdma, []).append(response[shortened_choice]['score'] / attribute.factor) attribute_prediction_reasonings.setdefault(choice, {}) # Choice level reasoning try: attribute_prediction_reasonings[choice].setdefault( - attribute.kdma, []).append(response[choice]['reasoning']) + attribute.kdma, []).append(response[shortened_choice]['reasoning']) # Probe level reasoning except KeyError: attribute_prediction_reasonings[choice].setdefault( diff --git a/align_system/algorithms/direct_regression_adm_component.py b/align_system/algorithms/direct_regression_adm_component.py index 3affcf79..f0146d8b 100644 --- a/align_system/algorithms/direct_regression_adm_component.py +++ b/align_system/algorithms/direct_regression_adm_component.py @@ -234,3 +234,120 @@ def _generic_object_repr(obj): num_samples={self.num_samples}, target_attribute_names_override={self.target_attribute_names_override}, )""", flags=re.MULTILINE).strip() + + +class OWDirectRegressionADMComponent(DirectRegressionADMComponent): + def run(self, + scenario_state, + choices, + icl_dialog_elements=[], + alignment_target=None): + if alignment_target is None: + target_attribute_names = [] + else: + target_attribute_names = attributes_in_alignment_target(alignment_target) + + if self.target_attribute_names_override is not None: + overridden_target_attribute_names = [] + for attribute_name in self.target_attribute_names_override: + if attribute_name == '*': + # '*' in the override means to include the attribute names + # from the target (in addition to whatever else is + # specified in the override) + overridden_target_attribute_names.extend(target_attribute_names) + else: + overridden_target_attribute_names.append(attribute_name) + + target_attribute_names = overridden_target_attribute_names + + if self.enable_caching: + scenario_state_copy = copy.deepcopy(scenario_state) + if hasattr(scenario_state, 'elapsed_time'): + # Don't consider the elapsed_time of the state when caching + scenario_state_copy.elapsed_time = 0 + + depends = '\n'.join(( + self.cache_repr(), + repr(scenario_state_copy), + repr(choices), + repr(icl_dialog_elements), + repr(target_attribute_names))) + + cacher = ub.Cacher('direct_regression_adm_component', depends, verbose=0) + log.debug(f'cacher.fpath={cacher.fpath}') + + cached_output = cacher.tryload() + if cached_output is not None: + log.info("Cache hit for `direct_regression_adm_component`" + " returning cached output") + return cached_output + else: + log.info("Cache miss for `direct_regression_adm_component` ..") + + if not isinstance(scenario_state, Mapping): + scenario_state = scenario_state.to_dict() + + attribute_prediction_scores = {} + attribute_prediction_reasonings = {} + for attribute in target_attribute_names: + if attribute not in self.per_attribute_templates: + raise RuntimeError(f"Missing {attribute} from self.per_attribute_templates") + + for choice in choices: + dialog = [] + system_prompt = self.per_attribute_templates[attribute]['system_prompt'] + if callable(system_prompt): + system_prompt = call_with_coerced_args( + system_prompt, + {'model':self.structured_inference_engine.model} + ) + elif not isinstance(system_prompt, str): + raise RuntimeError("system_prompt is of an unexpected type") + + dialog.insert(0, DialogElement(role='system', content=system_prompt)) + + prompt_template = self.per_attribute_templates[attribute]['prompt_template'] + if callable(prompt_template): + prompt = call_with_coerced_args( + prompt_template, + {'choice': choice, + 'scenario_state': scenario_state}) + elif isinstance(prompt_template, str): + prompt = Template(prompt_template).render( + {'choice': choice, + 'scenario_state': scenario_state}) + else: + raise RuntimeError("prompt_template is of an unexpected type") + + dialog.append(DialogElement(role='user', + content=prompt)) + + dialog_prompt = self.structured_inference_engine.dialog_to_prompt(dialog) + + log.info(f"[bold]*{attribute.upper()} PREDICTION DIALOG PROMPT*[/bold]", + extra={"markup": True}) + log.info(dialog_prompt) + + output_schema = call_with_coerced_args( + self.per_attribute_templates[attribute]['schema_template'], {}) + + responses = self.structured_inference_engine.run_inference( + [dialog_prompt] * self.num_samples, output_schema) + + for i, response in enumerate(responses): + log.info(f"[bold]*{attribute.upper()} PREDICTION RESPONSE (sample #{i})*[/bold]", extra={"markup": True}) + log.info(response, extra={"highlighter": JSON_HIGHLIGHTER}) + + factor = self.per_attribute_templates[attribute].get('factor', 100) + attribute_prediction_scores.setdefault(choice, {})[attribute] =\ + [r['score'] / float(factor) for r in responses] + attribute_prediction_reasonings.setdefault(choice, {})[attribute] =\ + [r['reasoning'] for r in responses] + + outputs = (attribute_prediction_reasonings, + attribute_prediction_scores) + + if self.enable_caching: + cacher.save(outputs) + + return outputs diff --git a/align_system/algorithms/open_world_components.py b/align_system/algorithms/open_world_components.py index 5095cd97..1fc6f4b0 100644 --- a/align_system/algorithms/open_world_components.py +++ b/align_system/algorithms/open_world_components.py @@ -1,5 +1,9 @@ +import copy +import inspect import json +import re from collections import defaultdict +import ubelt as ub from rich.highlighter import JSONHighlighter from swagger_client.models import ActionTypeEnum, CharacterTagEnum @@ -7,8 +11,17 @@ from align_system.algorithms.outlines_baseline_adm_component import OutlinesBaselineADMComponent from align_system.algorithms.alignment_adm_component import MedicalOnlyAlignmentADMComponent from align_system.data_models.dialog import DialogElement -from align_system.prompt_engineering.outlines_prompts import character_choice_json_schema, tag_choice_json_schema -from align_system.prompt_engineering.ow_prompts import FollowupClarifyCharacterPrompt, FollowupClarifyTagPrompt +from align_system.prompt_engineering.outlines_prompts import ( + character_choice_json_schema, + tag_choice_json_schema, + treatment_choice_from_list_json_schema +) +from align_system.prompt_engineering.ow_prompts import ( + FollowupClarifyCharacterPrompt, + FollowupClarifyTagPrompt, + FollowupClarifyTreatmentPrompt, + OWPart3CharacterDescriptionWVitals +) from align_system.utils import call_with_coerced_args, logging, get_swagger_class_enum_values log = logging.getLogger(__name__) @@ -16,6 +29,9 @@ class OWFormatChoicesADMComponent(ADMComponent): + def __init__(self): + self.choice_template = OWPart3CharacterDescriptionWVitals() + def run_returns(self): return ('choices', 'choice_to_action_mapping') @@ -28,7 +44,7 @@ def run(self, scenario_state, actions): ] character_to_choice = { - c.id: f"{c.name}: {c.unstructured}" + c.id: self.choice_template(c).rstrip() for c in scenario_state.characters } @@ -74,8 +90,25 @@ def run( elif len(possible_actions) == 1: # Single action, choose that chosen_action = possible_actions[0] else: + relevant_char_ids = set() + filter_by_ids = True + for a in possible_actions: + if a.character_id is not None: + relevant_char_ids.add(a.character_id) + else: # Could be any character + filter_by_ids = False + break + + filtered_scenario_state = copy.deepcopy(scenario_state) + if filter_by_ids: + filtered_scenario_state_characters = [] + for c in scenario_state.characters: + if c.id in relevant_char_ids: + filtered_scenario_state_characters.append(c) + filtered_scenario_state.characters = filtered_scenario_state_characters + choices = [a.unstructured for a in possible_actions] - chosen_choice, justification, choice_to_action_dialog = super().run(scenario_state, choices) + chosen_choice, justification, choice_to_action_dialog = super().run(filtered_scenario_state, choices) chosen_action = possible_actions[choices.index(chosen_choice)] @@ -96,13 +129,16 @@ def __init__( structured_inference_engine, scenario_description_template, system_prompt=None, + enable_caching=False, ): self.structured_inference_engine = structured_inference_engine self.scenario_description_template = scenario_description_template self.system_prompt = system_prompt + self.enable_caching = enable_caching self.followup_character_prompt = FollowupClarifyCharacterPrompt() self.followup_tag_prompt = FollowupClarifyTagPrompt() + self.followup_treatment_prompt = FollowupClarifyTreatmentPrompt() def run_returns(self): return ('chosen_action', 'action_parameter_completion_dialog') @@ -112,6 +148,30 @@ def run( scenario_state, chosen_action, ): + if self.enable_caching: + scenario_state_copy = copy.deepcopy(scenario_state) + if hasattr(scenario_state, 'elapsed_time'): + # Don't consider the elapsed_time of the state when caching + scenario_state_copy.elapsed_time = 0 + + depends = '\n'.join(( + self.cache_repr(), + repr(scenario_state_copy), + repr(chosen_action))) + + cacher = ub.Cacher('ow_action_parameter_completion_adm_component', depends, verbose=0) + log.debug(f'cacher.fpath={cacher.fpath}') + + cached_output = cacher.tryload() + if cached_output is not None: + log.info("Cache hit for `ow_action_parameter_completion_adm_component`" + " returning cached output") + return cached_output + else: + log.info("Cache miss for `ow_action_parameter_completion_adm_component` ..") + + # Make a deepcopy of chosen_action so in-place modifications don't mutate input state if un-cached + chosen_action = copy.deepcopy(chosen_action) action_parameter_completion_dialog = {} # Action requires a character ID @@ -200,13 +260,95 @@ def run( action_parameter_completion_dialog["tag"] = dialog - return chosen_action, action_parameter_completion_dialog + # Treatment requires a selected treatment supply + if chosen_action.action_type == ActionTypeEnum.TREAT_PATIENT: + dialog = [] + if self.system_prompt is not None: + dialog.append(DialogElement(role='system', content=self.system_prompt())) + + if chosen_action.parameters is None: + chosen_action.parameters = {} + + if 'treatment' not in chosen_action.parameters: + chosen_character = None + for c in scenario_state.characters: + if c.id == chosen_action.character_id: + chosen_character = c + break + + supplies_dict = {s.type.value: s.quantity for s in scenario_state.supplies + if s.quantity > 0} + # TODO: Better handle this corner case + assert len(supplies_dict) > 0 + dialog.append( + DialogElement(role='user', content=self.followup_treatment_prompt(chosen_character, supplies_dict)) + ) + + dialog_prompt = self.structured_inference_engine.dialog_to_prompt(dialog) + log.info("[bold]*TREATMENT FOLLOWUP PROMPT*[/bold]", extra={"markup": True}) + log.info(dialog_prompt) + + valid_treatments = list(supplies_dict.keys()) + selected_treatment = self.structured_inference_engine.run_inference( + dialog_prompt, + treatment_choice_from_list_json_schema(json.dumps(valid_treatments)) + ) + log.info("[bold]*TREATMENT FOLLOWUP RESPONSE*[/bold]", extra={"markup": True}) + log.info(selected_treatment, extra={"highlighter": JSON_HIGHLIGHTER}) + + chosen_action.parameters['treatment'] = selected_treatment["treatment_choice"] + + justification = selected_treatment["brief_reasoning"] + if isinstance(chosen_action, tuple) and hasattr(chosen_action, "_replace"): + chosen_action = chosen_action._replace(justification=justification) + else: + chosen_action.justification = justification + + action_parameter_completion_dialog["treatment"] = dialog + + outputs = (chosen_action, action_parameter_completion_dialog) + + if self.enable_caching: + cacher.save(outputs) + + return outputs + + def cache_repr(self): + ''' + Return a string representation of this object for caching; + .i.e. if the return value of this function is the same for two + object instances, it's assumed that `run` output will be + the same if given the same parameters + ''' + + def _generic_object_repr(obj): + if obj is None: + return "None" + + init_params = inspect.signature(obj.__class__.__init__).parameters + obj_vars = vars(obj) + + return "{}.{}({})".format( + obj.__class__.__module__, + obj.__class__.__name__, + ", ".join([f"{p}={obj_vars[p]}" for p in init_params + if p != 'self' and p != 'args' and p != 'kwargs'])) + + return re.sub(r'^\s+', '', + f""" + {self.__class__.__module__}.{self.__class__.__name__}( + structured_inference_engine={self.structured_inference_engine.cache_repr()}, + scenario_description_template={_generic_object_repr(self.scenario_description_template)}, + system_prompt={_generic_object_repr(self.system_prompt) if callable(self.system_prompt) else self.system_prompt}, + )""", flags=re.MULTILINE).strip() + class OWTaggingAdjustmentADMComponent(MedicalOnlyAlignmentADMComponent): def __init__(self, increase_priority_threshold=0.55, decrease_priority_threshold=0.80): self.increase_priority_threshold = increase_priority_threshold self.decrease_priority_threshold = decrease_priority_threshold + def run_returns(self): return ('chosen_action') @@ -251,10 +393,10 @@ def _get_sorted_ranking(ratings, descending=True): ranking_delta = medical_ranking - aligned_ranking percent_change = ranking_delta / len(choices) - # What tags have been given out already + # What tags have been given out already (nearby) tag_counts = defaultdict(int) for c in scenario_state.characters: - if c.tag is not None: + if c.tag is not None and not c.unseen and c.nearby: tag_counts[c.tag] += 1 lowest_priority_given_idx = None for i in range(len(tag_order)-2, -1, -1): # Don't consider black tags, order is slightly weird @@ -267,7 +409,7 @@ def _get_sorted_ranking(ratings, descending=True): adjusted_tag_idx = max(0, adjusted_tag_idx - 1) if percent_change < -self.decrease_priority_threshold: adjusted_tag_idx = min(len(tag_order)-1, adjusted_tag_idx + 1) - # Have already given out lower priority tags, heuristic only works when omniscient + # Check if we have already given out lower priority tags nearby if lowest_priority_given_idx is not None and lowest_priority_given_idx > assigned_tag_idx: adjusted_tag_idx = lowest_priority_given_idx @@ -282,3 +424,11 @@ def _get_sorted_ranking(ratings, descending=True): ) return chosen_action + + +class OWChoiceSchemaTransform: + def __call__(self, choice): + if ":" in choice: + return choice.split(':', 1)[0] + else: + return choice diff --git a/align_system/algorithms/pipeline_adm.py b/align_system/algorithms/pipeline_adm.py index 95b56a46..7813948c 100644 --- a/align_system/algorithms/pipeline_adm.py +++ b/align_system/algorithms/pipeline_adm.py @@ -11,6 +11,11 @@ class PipelineADM(ActionBasedADM): def __init__(self, steps: list[ADMComponent]): self.steps = steps + def reset_history(self): + for step in self.steps: + if hasattr(step, 'reset_history'): + step.reset_history() + def choose_action(self, scenario_state, available_actions, diff --git a/align_system/algorithms/random_adm_component.py b/align_system/algorithms/random_adm_component.py index 7c3b5cf0..f58e146f 100644 --- a/align_system/algorithms/random_adm_component.py +++ b/align_system/algorithms/random_adm_component.py @@ -99,12 +99,23 @@ def run(self, # Action requires a character ID if chosen_action.action_type in {'TREAT_PATIENT', ActionTypeEnum.MOVE_TO_EVAC, - ActionTypeEnum.TAG_CHARACTER}: + ActionTypeEnum.TAG_CHARACTER, + 'CHECK_VITALS'}: + if chosen_action.character_id is None: + chosen_action.character_id = random.choice([ + c.id + for c in scenario_state.characters + if hasattr(c, "unseen") and not c.unseen + if hasattr(c, "nearby") and c.nearby + ]) + + elif chosen_action.action_type == 'MOVE_TO': if chosen_action.character_id is None: chosen_action.character_id = random.choice([ c.id for c in scenario_state.characters if hasattr(c, "unseen") and not c.unseen + if hasattr(c, "nearby") and not c.nearby ]) if chosen_action.action_type == ActionTypeEnum.TAG_CHARACTER: @@ -115,6 +126,14 @@ def run(self, chosen_action.parameters['category'] = random.choice( get_swagger_class_enum_values(CharacterTagEnum)) + if chosen_action.action_type == 'TREAT_PATIENT': + if chosen_action.parameters is None: + chosen_action.parameters = {} + + chosen_action.parameters['treatment'] = random.choice([ + s.type.value for s in scenario_state.supplies + if s.quantity > 0]) + chosen_action.justification = "Random choice" return chosen_action diff --git a/align_system/algorithms/world_state_adm_component.py b/align_system/algorithms/world_state_adm_component.py new file mode 100644 index 00000000..069dbca2 --- /dev/null +++ b/align_system/algorithms/world_state_adm_component.py @@ -0,0 +1,206 @@ +from copy import deepcopy + +from align_system.algorithms.abstracts import ADMComponent +from align_system.utils import logging + +log = logging.getLogger(__name__) + + +def _get_field(obj, field, default=None): + if isinstance(obj, dict): + return obj.get(field, default) + return getattr(obj, field, default) + + +def _set_field(obj, field, value): + if isinstance(obj, dict): + obj[field] = value + else: + setattr(obj, field, value) + + +def _get_items(obj): + if isinstance(obj, dict): + return list(obj.items()) + elif hasattr(obj, 'to_dict'): + return list(obj.to_dict().items()) + elif hasattr(obj, '__dict__'): + return list(obj.__dict__.items()) + return [] + + +def _merge_dict_or_obj(existing_obj, new_obj, preserve_if_new_none=False): + """ + Recursively merges new_obj into existing_obj. + Supports both dicts and objects with properties/attributes. + """ + if new_obj is None: + return existing_obj + + if existing_obj is None: + return new_obj + + for key, new_val in _get_items(new_obj): + if new_val is None and preserve_if_new_none: + continue + + existing_val = _get_field(existing_obj, key) + + if (isinstance(new_val, dict) or hasattr(new_val, 'to_dict')) and not isinstance(new_val, list): + if existing_val is not None: + merged_sub = _merge_dict_or_obj(existing_val, new_val, preserve_if_new_none=preserve_if_new_none) + _set_field(existing_obj, key, merged_sub) + else: + _set_field(existing_obj, key, new_val) + else: + _set_field(existing_obj, key, new_val) + + return existing_obj + + +class WorldStateTrackerADMComponent(ADMComponent): + """ + ADM Component that maintains cumulative world state across scenes/probes in a scenario. + + Handles recursive merging of character details (e.g. `vitals`). Preserves detailed character + information acquired when `nearby == True` and vitals assessed during interaction, preventing + less-detailed descriptions from overwriting them when moving away (`nearby == False`). Dynamic + attributes (e.g. `tag`, `visited`, `unseen`, top-level `supplies`) continue to update normally. + """ + def __init__(self, + output_scenario_state_key='scenario_state', + output_original_state_key='original_scenario_state', + output_stale_info_key='world_state_stale_info'): + self.output_scenario_state_key = output_scenario_state_key + self.output_original_state_key = output_original_state_key + self.output_stale_info_key = output_stale_info_key + + self.current_scenario_id = None + self.world_state = None + self.stale_characters = {} + + def reset_history(self): + log.info("[bold]*Resetting WorldStateTrackerADMComponent history*[/bold]", + extra={"markup": True}) + self.current_scenario_id = None + self.world_state = None + self.stale_characters = {} + + def run_returns(self): + return (self.output_scenario_state_key, + self.output_original_state_key, + self.output_stale_info_key) + + def run(self, scenario_state, scenario_id=None): + # Auto-reset if scenario_id changes unexpectedly without an explicit reset_history call + if scenario_id is not None and scenario_id != self.current_scenario_id: + if self.current_scenario_id is not None: + log.info(f"Scenario ID changed from {self.current_scenario_id} to {scenario_id}, resetting world state") + self.reset_history() + self.current_scenario_id = scenario_id + + original_scenario_state = scenario_state + # Deepcopy the incoming scenario_state up front to ensure complete reference independence + new_state = deepcopy(scenario_state) + + if self.world_state is None: + # First scene of the scenario: initialize world state as deep copy of incoming state + self.world_state = new_state + self.stale_characters = {} + else: + # Merge incoming state into cumulative world_state + self._merge_scenario_state(new_state) + + stale_info = { + 'stale_characters': deepcopy(self.stale_characters) + } + + return self.world_state, original_scenario_state, stale_info + + def _merge_scenario_state(self, new_state): + # 1. Update top-level fields (unstructured narrative, elapsed_time, meta_info, supplies, etc.) + for key, val in _get_items(new_state): + if key != 'characters': + _set_field(self.world_state, key, val) + + # 2. Merge characters + existing_chars = _get_field(self.world_state, 'characters') + if existing_chars is None: + existing_chars = [] + _set_field(self.world_state, 'characters', existing_chars) + + existing_char_map = {} + for char in existing_chars: + cid = _get_field(char, 'id') or _get_field(char, 'name') + if cid: + existing_char_map[cid] = char + + new_chars = _get_field(new_state, 'characters') or [] + + for new_char in new_chars: + cid = _get_field(new_char, 'id') or _get_field(new_char, 'name') + if not cid: + existing_chars.append(new_char) + continue + + if cid not in existing_char_map: + existing_chars.append(new_char) + existing_char_map[cid] = new_char + else: + existing_char = existing_char_map[cid] + self._merge_character(existing_char, new_char, cid) + + def _merge_character(self, existing_char, new_char, cid): + old_nearby = bool(_get_field(existing_char, 'nearby', False)) + new_nearby = bool(_get_field(new_char, 'nearby', False)) + stale_fields_for_char = [] + + if old_nearby and not new_nearby: + # Character was previously nearby (rich detail) but is now not nearby + old_unstructured = _get_field(existing_char, 'unstructured') + old_vitals = deepcopy(_get_field(existing_char, 'vitals')) + + # Update all fields from new_char + for key, val in _get_items(new_char): + if key not in {'unstructured', 'vitals'}: + _set_field(existing_char, key, val) + + # Preserve richer unstructured description + if old_unstructured is not None: + _set_field(existing_char, 'unstructured', old_unstructured) + stale_fields_for_char.append('unstructured') + + # Merge/preserve vitals non-destructively + new_vitals = _get_field(new_char, 'vitals') + if old_vitals is not None: + if new_vitals is None: + _set_field(existing_char, 'vitals', old_vitals) + else: + merged_vitals = _merge_dict_or_obj(old_vitals, new_vitals, preserve_if_new_none=True) + _set_field(existing_char, 'vitals', merged_vitals) + stale_fields_for_char.append('vitals') + elif new_vitals is not None: + _set_field(existing_char, 'vitals', new_vitals) + + if stale_fields_for_char: + self.stale_characters[cid] = { + 'stale_fields': stale_fields_for_char, + 'nearby': False + } + log.debug(f"Preserving detailed info ({stale_fields_for_char}) for non-nearby character '{cid}'") + else: + # Character is now nearby (or was never nearby) -> recursive update + old_vitals = _get_field(existing_char, 'vitals') + + for key, val in _get_items(new_char): + if key == 'vitals': + if old_vitals is not None and val is not None: + merged_vitals = _merge_dict_or_obj(old_vitals, val) + _set_field(existing_char, key, merged_vitals) + elif val is not None: + _set_field(existing_char, key, val) + else: + _set_field(existing_char, key, val) + + if cid in self.stale_characters: + del self.stale_characters[cid] diff --git a/align_system/configs/adm_component/misc/world_state_tracker.yaml b/align_system/configs/adm_component/misc/world_state_tracker.yaml new file mode 100644 index 00000000..d538a008 --- /dev/null +++ b/align_system/configs/adm_component/misc/world_state_tracker.yaml @@ -0,0 +1 @@ +_target_: align_system.algorithms.world_state_adm_component.WorldStateTrackerADMComponent diff --git a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml index 7ba4ef20..907eef03 100644 --- a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml +++ b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml @@ -59,3 +59,5 @@ driver: force_determinism: true align_to_target: true save_last_unstructured_state_per_scenario: true +# CACI security policy flags our justifications for Feb/April scenarios +remove_justifications: true diff --git a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live_part1.yaml b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live_part1.yaml index 2d8df307..81063d57 100644 --- a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live_part1.yaml +++ b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_midpoint_live_part1.yaml @@ -34,3 +34,5 @@ driver: force_determinism: true align_to_target: true save_last_unstructured_state_per_scenario: true +# CACI security policy flags our justifications for Feb/April scenarios +remove_justifications: true diff --git a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_random_effects_live.yaml b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_random_effects_live.yaml index 71275ba6..5cf039ec 100644 --- a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_random_effects_live.yaml +++ b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_fewshot_comparative_regression_random_effects_live.yaml @@ -58,3 +58,5 @@ driver: force_determinism: true align_to_target: true save_last_unstructured_state_per_scenario: true +# CACI security policy flags our justifications for Feb/April scenarios +remove_justifications: true diff --git a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml index a7736bda..1539ff91 100644 --- a/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml +++ b/align_system/configs/experiment/phase2_feb_openworld/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml @@ -15,7 +15,7 @@ interface: api_endpoint: "https://darpaitm.caci.com" session_type: eval training_session: null - username: "ALIGN-ADM-Ph2-ComparativeRegression-Zeroshot-Mistral-7B-Instruct-v0.3" + username: "ALIGN-ADM-Ph2-Zeroshot-ComparativeRegression-Mistral-7B-Instruct-v0.3" domain: "p2triage" adm_profile: FEB_OPENWORLD2 diff --git a/align_system/configs/experiment/phase2_openworld_part3/baseline_live_test.yaml b/align_system/configs/experiment/phase2_openworld_part3/baseline_live_test.yaml new file mode 100644 index 00000000..dc8e5438 --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/baseline_live_test.yaml @@ -0,0 +1,55 @@ +# @package _global_ +defaults: + - override /adm: pipeline_baseline + - override /inference_engine@adm.structured_inference_engine: outlines_structured_greedy + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - override /interface: ta3 + - override /driver: itm_phase2_ow + +interface: + api_endpoint: "https://darpaitm.caci.com" +# api_endpoint: 'http://127.0.0.1:8081' + session_type: eval + training_session: null + username: "testrun-ALIGN-ADM-OutlinesBaseline-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + outlines_baseline: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + prompt_template: + _target_: align_system.prompt_engineering.outlines_prompts.Phase2BaselinePrompt + + enable_caching: true + + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.outlines_baseline} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: false + +save_last_unstructured_state_per_scenario: true + +hydra: + run: + dir: 'ow_part3_live_test/baseline/${now:%Y-%m-%d__%H-%M-%S}' diff --git a/align_system/configs/experiment/phase2_openworld_part3/baseline_with_world_state_live_test.yaml b/align_system/configs/experiment/phase2_openworld_part3/baseline_with_world_state_live_test.yaml new file mode 100644 index 00000000..12a0395c --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/baseline_with_world_state_live_test.yaml @@ -0,0 +1,57 @@ +# @package _global_ +defaults: + - /adm_component/misc@adm.step_definitions.world_state_tracker: world_state_tracker + - override /adm: pipeline_baseline + - override /inference_engine@adm.structured_inference_engine: outlines_structured_greedy + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - override /interface: ta3 + - override /driver: itm_phase2_ow + +interface: + api_endpoint: "https://darpaitm.caci.com" +# api_endpoint: 'http://127.0.0.1:8081' + session_type: eval + training_session: null + username: "testrun-ALIGN-ADM-OutlinesBaseline-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + outlines_baseline: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + prompt_template: + _target_: align_system.prompt_engineering.outlines_prompts.Phase2BaselinePrompt + + enable_caching: true + + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.world_state_tracker} + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.outlines_baseline} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: false + +save_last_unstructured_state_per_scenario: true + +hydra: + run: + dir: 'ow_part3_live_test/baseline-worldstate/${now:%Y-%m-%d__%H-%M-%S}' diff --git a/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_direct_regression_midpoint_live_test.yaml b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_direct_regression_midpoint_live_test.yaml new file mode 100644 index 00000000..3f55482d --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_direct_regression_midpoint_live_test.yaml @@ -0,0 +1,78 @@ +# @package _global_ +defaults: + - /adm_component/misc@adm.step_definitions.world_state_tracker: world_state_tracker + - /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - /adm_component/misc@adm.step_definitions.tagging_adjustment: ow_tagging_adjustment + - override /adm: phase2_pipeline_direct_regression + - override /interface: ta3 + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/alignment@adm.step_definitions.scalar_alignment: multinomial_medical_urgency_weighted_scalar + - override /adm_component/misc@adm.step_definitions.ensure_chosen_action: ow_choice_to_action + - override /driver: itm_phase2_ow + +interface: + api_endpoint: 'http://127.0.0.1:8081' + #api_endpoint: "https://darpaitm.caci.com" + session_type: eval + training_session: null + username: "testrun-ALIGN-ADM-Ph2-DirectRegression-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + direct_regression: + _target_: align_system.algorithms.direct_regression_adm_component.OWDirectRegressionADMComponent + enable_caching: true + + per_attribute_templates: + # Include vitals information for Medical Urgency + medical: + prompt_template: | + Given the injury described below, provide a reasoning and a medical urgency rating: + + {{ choice }} + + affiliation: + prompt_template: | + Given the casualty described below, provide a reasoning and an affiliation score: + + {{ choice }} + + merit: + prompt_template: | + Given the casualty described below, provide a reasoning and a merit score: + + {{ choice }} + + ensure_chosen_action: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.world_state_tracker} + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.direct_regression} + #- ${ref:adm.step_definitions.regression_rule_based_correction} + - ${ref:adm.step_definitions.scalar_alignment} + - ${ref:adm.step_definitions.justification_from_reasonings} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.tagging_adjustment} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: true +save_last_unstructured_state_per_scenario: true diff --git a/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml new file mode 100644 index 00000000..1225b83a --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_fewshot_comparative_regression_midpoint_live.yaml @@ -0,0 +1,77 @@ +# @package _global_ +defaults: + - /adm_component/misc@adm.step_definitions.world_state_tracker: world_state_tracker + - /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - /adm_component/misc@adm.step_definitions.tagging_adjustment: ow_tagging_adjustment + - override /adm: phase2_pipeline_fewshot_comparative_regression + - override /interface: ta3 + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/alignment@adm.step_definitions.scalar_alignment: multinomial_medical_urgency_weighted_scalar + - override /adm_component/misc@adm.step_definitions.ensure_chosen_action: ow_choice_to_action + - override /template/scenario_description@adm.scenario_description_template: phase2 + - override /driver: itm_phase2_ow + +interface: + api_endpoint: "https://darpaitm.caci.com" + session_type: eval + training_session: null + username: "ALIGN-ADM-Ph2-ComparativeRegression-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + regression_icl: + icl_generator_partial: + _target_: align_system.utils.incontext_utils.Phase2ComparativeRegressionIncontextExampleGeneratorOWConversion + incontext_settings: + number: 20 + datasets: + medical: /data/shared/samba/phase2_icl/Feb2026-MU-train_20251218.json + affiliation: /data/shared/samba/phase2_icl/Feb2026-AF-train_20251218.json + merit: /data/shared/samba/phase2_icl/Feb2026-MF-train_20251218.json + personal_safety: /data/shared/samba/phase2_icl/Feb2026-PS-train_20251218.json + search: /data/shared/samba/phase2_icl/Feb2026-SS-train_20251218.json + character_choice_template: + _target_: align_system.prompt_engineering.ow_prompts.OWPart3CharacterDescriptionWVitals + choice_schema_transform: + _target_: align_system.algorithms.open_world_components.OWChoiceSchemaTransform + enable_caching: true + comparative_regression: + enable_caching: true + choice_schema_transform: + _target_: align_system.algorithms.open_world_components.OWChoiceSchemaTransform + ensure_chosen_action: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.world_state_tracker} + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.regression_icl} + - ${ref:adm.step_definitions.comparative_regression} + #- ${ref:adm.step_definitions.regression_rule_based_correction} + - ${ref:adm.step_definitions.scalar_alignment} + - ${ref:adm.step_definitions.justification_from_reasonings} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.tagging_adjustment} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: true +save_last_unstructured_state_per_scenario: true +# CACI security policy flags our justifications for Feb/April scenarios +remove_justifications: true diff --git a/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_fewshot_comparative_regression_midpoint_live_test.yaml b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_fewshot_comparative_regression_midpoint_live_test.yaml new file mode 100644 index 00000000..53514813 --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_fewshot_comparative_regression_midpoint_live_test.yaml @@ -0,0 +1,78 @@ +# @package _global_ +defaults: + - /adm_component/misc@adm.step_definitions.world_state_tracker: world_state_tracker + - /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - /adm_component/misc@adm.step_definitions.tagging_adjustment: ow_tagging_adjustment + - override /adm: phase2_pipeline_fewshot_comparative_regression + - override /interface: ta3 + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/alignment@adm.step_definitions.scalar_alignment: multinomial_medical_urgency_weighted_scalar + - override /adm_component/misc@adm.step_definitions.ensure_chosen_action: ow_choice_to_action + - override /template/scenario_description@adm.scenario_description_template: phase2 + - override /driver: itm_phase2_ow + +interface: + api_endpoint: 'http://127.0.0.1:8081' + #api_endpoint: "https://darpaitm.caci.com" + session_type: eval + training_session: null + username: "testrun-ALIGN-ADM-Ph2-ComparativeRegression-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + regression_icl: + icl_generator_partial: + _target_: align_system.utils.incontext_utils.Phase2ComparativeRegressionIncontextExampleGeneratorOWConversion + incontext_settings: + number: 20 + datasets: + medical: /data/shared/samba/phase2_icl/Feb2026-MU-train_20251218.json + affiliation: /data/shared/samba/phase2_icl/Feb2026-AF-train_20251218.json + merit: /data/shared/samba/phase2_icl/Feb2026-MF-train_20251218.json + personal_safety: /data/shared/samba/phase2_icl/Feb2026-PS-train_20251218.json + search: /data/shared/samba/phase2_icl/Feb2026-SS-train_20251218.json + character_choice_template: + _target_: align_system.prompt_engineering.ow_prompts.OWPart3CharacterDescriptionWVitals + choice_schema_transform: + _target_: align_system.algorithms.open_world_components.OWChoiceSchemaTransform + enable_caching: true + comparative_regression: + enable_caching: true + choice_schema_transform: + _target_: align_system.algorithms.open_world_components.OWChoiceSchemaTransform + ensure_chosen_action: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.world_state_tracker} + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.regression_icl} + - ${ref:adm.step_definitions.comparative_regression} + #- ${ref:adm.step_definitions.regression_rule_based_correction} + - ${ref:adm.step_definitions.scalar_alignment} + - ${ref:adm.step_definitions.justification_from_reasonings} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.tagging_adjustment} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: true +save_last_unstructured_state_per_scenario: true +# CACI security policy flags our justifications for Feb/April scenarios +remove_justifications: true diff --git a/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml new file mode 100644 index 00000000..5cb1082b --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_zeroshot_comparative_regression_midpoint_live.yaml @@ -0,0 +1,58 @@ +# @package _global_ +defaults: + - /adm_component/misc@adm.step_definitions.world_state_tracker: world_state_tracker + - /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - /adm_component/misc@adm.step_definitions.tagging_adjustment: ow_tagging_adjustment + - override /adm: phase2_pipeline_zeroshot_comparative_regression + - override /interface: ta3 + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/alignment@adm.step_definitions.scalar_alignment: multinomial_medical_urgency_weighted_scalar + - override /adm_component/misc@adm.step_definitions.ensure_chosen_action: ow_choice_to_action + - override /template/scenario_description@adm.scenario_description_template: phase2 + - override /driver: itm_phase2_ow + +interface: + api_endpoint: "https://darpaitm.caci.com" + session_type: eval + training_session: null + username: "ALIGN-ADM-Ph2-Zeroshot-ComparativeRegression-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + comparative_regression: + enable_caching: true + choice_schema_transform: + _target_: align_system.algorithms.open_world_components.OWChoiceSchemaTransform + ensure_chosen_action: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.world_state_tracker} + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.comparative_regression} + #- ${ref:adm.step_definitions.regression_rule_based_correction} + - ${ref:adm.step_definitions.scalar_alignment} + - ${ref:adm.step_definitions.justification_from_reasonings} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.tagging_adjustment} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: true +save_last_unstructured_state_per_scenario: true diff --git a/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_zeroshot_comparative_regression_midpoint_live_test.yaml b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_zeroshot_comparative_regression_midpoint_live_test.yaml new file mode 100644 index 00000000..e15562a7 --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/phase2_pipeline_zeroshot_comparative_regression_midpoint_live_test.yaml @@ -0,0 +1,59 @@ +# @package _global_ +defaults: + - /adm_component/misc@adm.step_definitions.world_state_tracker: world_state_tracker + - /adm_component/misc@adm.step_definitions.action_parameter_completion: ow_action_parameter_completion + - /adm_component/misc@adm.step_definitions.tagging_adjustment: ow_tagging_adjustment + - override /adm: phase2_pipeline_zeroshot_comparative_regression + - override /interface: ta3 + - override /adm_component/misc@adm.step_definitions.format_choices: ow_format_choices + - override /adm_component/alignment@adm.step_definitions.scalar_alignment: multinomial_medical_urgency_weighted_scalar + - override /adm_component/misc@adm.step_definitions.ensure_chosen_action: ow_choice_to_action + - override /template/scenario_description@adm.scenario_description_template: phase2 + - override /driver: itm_phase2_ow + +interface: + api_endpoint: 'http://127.0.0.1:8081' + #api_endpoint: "https://darpaitm.caci.com" + session_type: eval + training_session: null + username: "testrun-ALIGN-ADM-Ph2-Zeroshot-ComparativeRegression-Mistral-7B-Instruct-v0.3" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + +adm: + step_definitions: + comparative_regression: + enable_caching: true + choice_schema_transform: + _target_: align_system.algorithms.open_world_components.OWChoiceSchemaTransform + ensure_chosen_action: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + action_parameter_completion: + scenario_description_template: + _target_: align_system.prompt_engineering.outlines_prompts.OWPart3ScenarioDescriptionWVitals + enable_caching: true + + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.world_state_tracker} + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.comparative_regression} + #- ${ref:adm.step_definitions.regression_rule_based_correction} + - ${ref:adm.step_definitions.scalar_alignment} + - ${ref:adm.step_definitions.justification_from_reasonings} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.tagging_adjustment} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: false + apply_action_filtering: true + +force_determinism: true +align_to_target: true +save_last_unstructured_state_per_scenario: true diff --git a/align_system/configs/experiment/phase2_openworld_part3/random_live_test.yaml b/align_system/configs/experiment/phase2_openworld_part3/random_live_test.yaml new file mode 100644 index 00000000..e3fdcabd --- /dev/null +++ b/align_system/configs/experiment/phase2_openworld_part3/random_live_test.yaml @@ -0,0 +1,37 @@ +# @package _global_ +defaults: + - override /adm: pipeline_random + - override /interface: ta3 + - override /driver: itm_phase2_ow + - override /adm_component/misc@adm.step_definitions.random_action_parameter_completion: ow_random_action_parameter_completion + +interface: +# api_endpoint: "https://darpaitm.caci.com" + api_endpoint: 'http://127.0.0.1:8081' + session_type: eval + training_session: null + username: "testrun-ALIGN-ADM-Random" + domain: "owtriage" + adm_profile: FEB_OPENWORLD3 + + +adm: + instance: + steps: + # Reference the step instances we want to use in order + - ${ref:adm.step_definitions.format_choices} + - ${ref:adm.step_definitions.random_choice} + - ${ref:adm.step_definitions.random_action_parameter_completion} + # - ${ref:adm.step_definitions.action_parameter_completion} + - ${ref:adm.step_definitions.ensure_chosen_action} + - ${ref:adm.step_definitions.populate_choice_info} + +driver: + expand_actions: true + expand_tagging: true + apply_action_filtering: true + +force_determinism: true +align_to_target: false + +save_last_unstructured_state_per_scenario: true diff --git a/align_system/drivers/itm_open_world.py b/align_system/drivers/itm_open_world.py index a29ab531..464b2e6c 100644 --- a/align_system/drivers/itm_open_world.py +++ b/align_system/drivers/itm_open_world.py @@ -40,7 +40,12 @@ def _expand_action_by_character(self, action, characters): for character in characters: new_action = deepcopy(action) new_action.character_id = character.id - new_action.unstructured = re.sub(r"(a )?Patient", character.name, action.unstructured) + new_action.unstructured = re.sub( + r"((?:a|the) )?Patient", + character.name, + action.unstructured, + flags=re.IGNORECASE + ) expanded_actions.append(new_action) @@ -131,6 +136,9 @@ def _compute_time_stats(times_s): "raw_times_s": times_s } + # Tracking last action and state to help prevent ADMs getting stuck in a loop + last_action = None + last_state = None # Loop through available scenarios while scenario := interface.start_scenario(): if scenario.id() == '': @@ -176,7 +184,9 @@ def _compute_time_stats(times_s): last_scene_id = None treated_patients = set() + tagged_patients = set() evac_patients = set() + checked_patients = set() while not scenario_complete: current_scene_id = current_state.meta_info.scene_id @@ -229,6 +239,18 @@ def _compute_time_stats(times_s): characters=current_state.characters )) + elif a.action_type == ActionTypeEnum.MOVE_TO: + available_actions_expanded.extend(self._expand_action_by_character( + action=a, + characters=current_state.characters + )) + + elif a.action_type == ActionTypeEnum.CHECK_VITALS: + available_actions_expanded.extend(self._expand_action_by_character( + action=a, + characters=current_state.characters + )) + else: available_actions_expanded.append(a) @@ -243,13 +265,16 @@ def _compute_time_stats(times_s): available_actions_filtered = [] for a in available_actions_expanded: if a.action_type == ActionTypeEnum.END_SCENE: - # We want to restrict end scene until all characters have been treated - continue + if(len(treated_patients) < len(current_state.characters) or + len(tagged_patients) < len(current_state.characters)): + # We want to restrict end scene until all characters have been tagged and treated + continue elif a.action_type == ActionTypeEnum.TAG_CHARACTER: untagged_characters = { c.id for c in current_state.characters if c.tag is None and not c.unseen + and c.nearby } if len(untagged_characters) == 0: # No more patients to tag continue @@ -263,6 +288,7 @@ def _compute_time_stats(times_s): treatable_patients = { c.id for c in current_state.characters if c.id not in treated_patients + and c.nearby } if len(treatable_patients) == 0: # No more patients to treat continue @@ -279,6 +305,68 @@ def _compute_time_stats(times_s): if a.character_id is not None and a.character_id not in evacable_patients: continue + elif a.action_type == ActionTypeEnum.MOVE_TO: + distant_patients = { + c.id for c in current_state.characters + if not c.nearby + } + if len(distant_patients) == 0: # No more patients to move to + continue + if a.character_id is not None and a.character_id not in distant_patients: + continue + + nearby_unchecked_patients = { + c.id for c in current_state.characters + if c.nearby and c.id not in checked_patients} + distant_unchecked_patients = { + c.id for c in current_state.characters + if not c.nearby and c.id not in checked_patients} + + nearby_untagged_patients = { + c.id for c in current_state.characters + if c.nearby and c.id not in tagged_patients} + distant_untagged_patients = { + c.id for c in current_state.characters + if not c.nearby and c.id not in tagged_patients} + + nearby_untreated_patients = { + c.id for c in current_state.characters + if c.nearby and c.id not in treated_patients} + distant_untreated_patients = { + c.id for c in current_state.characters + if not c.nearby and c.id not in treated_patients} + + # More strict than the other MOVE_TO + # filters above, only allow MOVE_TO to an + # untreated, unchecked, or untagged + # patient + if a.character_id is not None and a.character_id not in set.union( + distant_untreated_patients, + distant_unchecked_patients, + distant_untagged_patients): + continue + + if last_action is not None and last_action.action_type == ActionTypeEnum.MOVE_TO: + # If last_action was MOVE_TO and there + # are non-MOVE_TO actions to take on + # nearby patients you cannot MOVE_TO + # again (until you take a non-MOVE_TO + # action) + if len(set.union(nearby_untreated_patients, + nearby_untagged_patients, + nearby_unchecked_patients)) > 0: + continue + + elif a.action_type == ActionTypeEnum.CHECK_VITALS: + nearby_unchecked_patients = { + c.id for c in current_state.characters + if c.nearby and c.id not in checked_patients + } + if len(nearby_unchecked_patients) == 0: # No more patients to check + continue + if a.character_id is not None and a.character_id not in nearby_unchecked_patients: + continue + available_actions_filtered.append(a) log.debug("[bold]*AVAILABLE ACTIONS FILTERED*[/bold]", @@ -296,7 +384,7 @@ def _compute_time_stats(times_s): if end_scene_idx is not None: log.info("** All patients have been tagged and treated, ending scene") - action_to_take = available_actions[end_scene_idx] + action_to_take = available_actions_expanded[end_scene_idx] action_to_take.justification = "All patients have been tagged and treated" else: raise RuntimeError("No available actions from filtered list!") @@ -381,7 +469,10 @@ def _compute_time_stats(times_s): with open(save_input_output_to_path, 'w') as f: json.dump(inputs_outputs, f, indent=2) + last_state = current_state try: + if cfg.get('remove_justifications', False): + action_to_take.justification = "Dummy justification" if hasattr(action_to_take, "intent_action") and action_to_take.intent_action: current_state = scenario.intend_action(action_to_take) else: @@ -393,9 +484,20 @@ def _compute_time_stats(times_s): # If we treated a patient, record that treatment so we can ensure we treat everyone if action_to_take.action_type == ActionTypeEnum.TREAT_PATIENT: treated_patients.add(action_to_take.character_id) + # If we tagged a patient, record that tagging so we can ensure we tag everyone + if action_to_take.action_type == ActionTypeEnum.TAG_CHARACTER: + tagged_patients.add(action_to_take.character_id) # If we evaced a patient, record that so we don't try to evac them again if action_to_take.action_type == ActionTypeEnum.MOVE_TO_EVAC: evac_patients.add(action_to_take.character_id) + # If we checked vitals on a patient, record that so we don't check them again + # TODO: If in the future it's possible for vitals to + # change as the scenario progresses then we need + # something more sophisticated here + if action_to_take.action_type == ActionTypeEnum.CHECK_VITALS: + checked_patients.add(action_to_take.character_id) + + last_action = action_to_take scenario_complete = current_state.scenario_complete @@ -420,6 +522,9 @@ def _compute_time_stats(times_s): with open(final_scenario_state_output_path, "w") as f: print(current_state.unstructured, file=f) + last_action = None + last_state = None + if save_timing_to_path is not None: action_times["scenarios"].append(_compute_time_stats(sce_times_s)) diff --git a/align_system/interfaces/ta3_caci_action_based_service.py b/align_system/interfaces/ta3_caci_action_based_service.py index 5c8eb565..89e79c50 100644 --- a/align_system/interfaces/ta3_caci_action_based_service.py +++ b/align_system/interfaces/ta3_caci_action_based_service.py @@ -13,6 +13,8 @@ log = logging.getLogger(__name__) +PHASE2_DOMAINS = {"p2triage", "owtriage"} + class TA3CACIActionBasedServiceInterface(Interface): def __init__(self, @@ -122,7 +124,7 @@ def _take_or_intend_action(self, action, take_or_intend): if isinstance(action, dict): action = Action(**action) - if self.domain == "p2triage": + if self.domain in PHASE2_DOMAINS: updated_state = take_or_intend( session_id=self.session_id, action=action) @@ -152,7 +154,7 @@ def get_state(self): state = self.connection.get_scenario_state( session_id=self.session_id, scenario_id=self.scenario.id) - if self.domain == "p2triage": + if self.domain in PHASE2_DOMAINS: if state.threat_state is not None: state.unstructured = "{}\n{}".format( state.threat_state.unstructured, diff --git a/align_system/prompt_engineering/outlines_prompts.py b/align_system/prompt_engineering/outlines_prompts.py index 080b9889..8951d5b2 100644 --- a/align_system/prompt_engineering/outlines_prompts.py +++ b/align_system/prompt_engineering/outlines_prompts.py @@ -1231,6 +1231,38 @@ def __call__(self, scenario_state): return phase2_scenario_state_description_w_casualty_info(scenario_state) +@compat_outlines_prompt +def ow_part3_scenario_state_description_w_vitals(scenario_state): + """ + {{ scenario_state.unstructured.rstrip() }} + + {% if scenario_state.characters|length > 1 %} + Casualties: + {% endif %} + {% for c in scenario_state.characters %} + {% if scenario_state.characters|length > 1 %}- {% endif %}{{ c.name }}: {{ c.unstructured }} + {% if c.vitals is none or (c.vitals.avpu is none and c.vitals.breathing is none and c.vitals.heart_rate is none) %} + Vitals: Unknown + {% else %} + Vitals: + {% if c.vitals.avpu is not none %} + - AVPU: {{ c.vitals.avpu.value if c.vitals.avpu.value is defined else c.vitals.avpu }} + {% endif %} + {% if c.vitals.breathing is not none %} + - Breathing: {{ c.vitals.breathing.value if c.vitals.breathing.value is defined else c.vitals.breathing }} + {% endif %} + {% if c.vitals.heart_rate is not none %} + - Heart Rate: {{ c.vitals.heart_rate.value if c.vitals.heart_rate.value is defined else c.vitals.heart_rate }} + {% endif %} + {% endif %} + {% endfor %} + """ + +class OWPart3ScenarioDescriptionWVitals(): + def __call__(self, scenario_state): + return ow_part3_scenario_state_description_w_vitals(scenario_state) + + @compat_outlines_prompt def phase2_baseline_prompt(scenario_description, choices): """ diff --git a/align_system/prompt_engineering/ow_prompts.py b/align_system/prompt_engineering/ow_prompts.py index 69f5e71c..418e24ba 100644 --- a/align_system/prompt_engineering/ow_prompts.py +++ b/align_system/prompt_engineering/ow_prompts.py @@ -6,7 +6,7 @@ @compat_outlines_prompt -def followup_clarify_tag(character): +def followup_clarify_tag(c): """ Given the description below, please state what triage category \ should be applied to the patient given their injuries and vitals. \ @@ -21,7 +21,21 @@ def followup_clarify_tag(character): resources during the incident. This category should be used sparingly. Description: - {{ character.unstructured }} + {{ c.unstructured }} + {% if c.vitals is none or (c.vitals.avpu is none and c.vitals.breathing is none and c.vitals.heart_rate is none) %} + Vitals: Unknown + {% else %} + Vitals: + {% if c.vitals.avpu is not none %} + - AVPU: {{ c.vitals.avpu.value if c.vitals.avpu.value is defined else c.vitals.avpu }} + {% endif %} + {% if c.vitals.breathing is not none %} + - Breathing: {{ c.vitals.breathing.value if c.vitals.breathing.value is defined else c.vitals.breathing }} + {% endif %} + {% if c.vitals.heart_rate is not none %} + - Heart Rate: {{ c.vitals.heart_rate.value if c.vitals.heart_rate.value is defined else c.vitals.heart_rate }} + {% endif %} + {% endif %} """ class FollowupClarifyTagPrompt: @@ -36,7 +50,7 @@ def followup_clarify_character(scenario_description, action): character should receive the action. Scenario: - {{ scenario_descriptoin }} + {{ scenario_description }} Selected action: {{ action.unstructured }} @@ -45,3 +59,62 @@ def followup_clarify_character(scenario_description, action): class FollowupClarifyCharacterPrompt: def __call__(self, scenario_description, action): return followup_clarify_character(scenario_description, action) + + +@compat_outlines_prompt +def followup_clarify_treatment(c, supplies): + """ + Given the description below, please state what treatment supply \ + should be applied to the patient given their injuries and vitals. \ + Select a treatment from the following list of available supplies \ + (quantities shown in parentheses): + {% for supply, qty in supplies.items() %} + - {{ supply }} (x{{ qty }}) + {% endfor %} + + Description: + {{ c.unstructured }} + {% if c.vitals is none or (c.vitals.avpu is none and c.vitals.breathing is none and c.vitals.heart_rate is none) %} + Vitals: Unknown + {% else %} + Vitals: + {% if c.vitals.avpu is not none %} + - AVPU: {{ c.vitals.avpu.value if c.vitals.avpu.value is defined else c.vitals.avpu }} + {% endif %} + {% if c.vitals.breathing is not none %} + - Breathing: {{ c.vitals.breathing.value if c.vitals.breathing.value is defined else c.vitals.breathing }} + {% endif %} + {% if c.vitals.heart_rate is not none %} + - Heart Rate: {{ c.vitals.heart_rate.value if c.vitals.heart_rate.value is defined else c.vitals.heart_rate }} + {% endif %} + {%- endif %} + """ + +class FollowupClarifyTreatmentPrompt: + def __call__(self, character, supplies): + return followup_clarify_treatment(character, supplies) + + +@compat_outlines_prompt +def ow_part3_character_description_w_vitals(c): + """ + {{ c.name }}: {{ c.unstructured }} + {% if c.vitals is none or (c.vitals.avpu is none and c.vitals.breathing is none and c.vitals.heart_rate is none) %} + Vitals: Unknown + {% else %} + Vitals: + {% if c.vitals.avpu is not none %} + - AVPU: {{ c.vitals.avpu.value if c.vitals.avpu.value is defined else c.vitals.avpu }} + {% endif %} + {% if c.vitals.breathing is not none %} + - Breathing: {{ c.vitals.breathing.value if c.vitals.breathing.value is defined else c.vitals.breathing }} + {% endif %} + {% if c.vitals.heart_rate is not none %} + - Heart Rate: {{ c.vitals.heart_rate.value if c.vitals.heart_rate.value is defined else c.vitals.heart_rate }} + {% endif %} + {%- endif %} + """ + +class OWPart3CharacterDescriptionWVitals: + def __call__(self, character): + return ow_part3_character_description_w_vitals(character) diff --git a/align_system/utils/adm_utils.py b/align_system/utils/adm_utils.py index a11dbad4..674761b1 100644 --- a/align_system/utils/adm_utils.py +++ b/align_system/utils/adm_utils.py @@ -22,8 +22,10 @@ def format_choices(choices, available_actions, scenario_state): # available actions as we'll use the selected index later to # map to the corresponding action choices = [] + for a in available_actions: - if(a.action_type == ActionTypeEnum.APPLY_TREATMENT + if((a.action_type == 'APPLY_TREATMENT' or + a.action_type == 'TREAT_PATIENT') and a.parameters is not None and len(a.parameters) > 0): choices.append(detailed_unstructured_treatment_action_text(a, character_id_to_name)) elif(a.action_type == ActionTypeEnum.TAG_CHARACTER diff --git a/align_system/utils/incontext_utils.py b/align_system/utils/incontext_utils.py index 1a38a703..24c71fcf 100644 --- a/align_system/utils/incontext_utils.py +++ b/align_system/utils/incontext_utils.py @@ -888,10 +888,26 @@ def get_chain_of_thought_reasoning(self, target_kdma, scores): cot_reasoning = f"{max_choice} demonstates {adjective} more {target_kdma['name']} than {min_choice}." return cot_reasoning + # TODO: Refactor IncontextExampleGenerators to take scenario # description and prompt templates as arguments and use # `call_with_coerced_args` class Phase2ComparativeRegressionIncontextExampleGeneratorOWConversion(Phase2ComparativeRegressionIncontextExampleGenerator): + def __init__( + self, + incontext_settings, + target_kdmas, + state_hydration_domain=None, + scenario_description_template=None, + scorer=None, + character_choice_template=None, + choice_schema_transform=None, + ): + self.character_choice_template = character_choice_template + self.choice_schema_transform = choice_schema_transform + + super().__init__(incontext_settings, target_kdmas, state_hydration_domain, scenario_description_template, scorer) + def set_icl_datasets(self): from swagger_client.models import ActionTypeEnum @@ -906,11 +922,12 @@ def set_icl_datasets(self): # Add each examples to icl_datasets for example in kdma_incontext_data: - character_unstructured_by_id = {c.id: c.unstructured for c in example['state'].characters} + character_by_id = {c.id: c for c in example['state'].characters} # Get example response icl_response = {} included_choices = [] + schema_choices = [] for action, choice, kdma_value in zip(example['actions'], example['choices'], example["kdma_values"]): # HACK: Reformat choice string for OW # TODO: Bring more in line with OWFormatChoicesADMComponent (align_system/algorithms/open_world_components.py) @@ -925,21 +942,29 @@ def set_icl_datasets(self): # expansion or?? assert c_id is not None - choice = f"{c_id}: {character_unstructured_by_id[c_id]}" + if self.character_choice_template is None: + choice = f"{c_id}: {character_by_id[c_id].unstructured}" + else: + choice = self.character_choice_template(character_by_id[c_id]).rstrip() # Only include choice if there is a ground truth KDMA value available if kdma_value is None: continue # Groundtruth KDMA values are 0-1, but ADM may predict on a different scale scaled_kdma_value = int(kdma_value * target_kdma["factor"]) - icl_response[choice] = {} - icl_response[choice]['score'] = scaled_kdma_value + schema_choice = choice + if self.choice_schema_transform is not None: + schema_choice = self.choice_schema_transform(schema_choice) + icl_response[schema_choice] = {} + icl_response[schema_choice]['score'] = scaled_kdma_value included_choices.append(choice) + schema_choices.append(schema_choice) icl_response_with_reasoning={} icl_response_with_reasoning['reasoning'] = self.get_chain_of_thought_reasoning(target_kdma, icl_response) icl_response_with_reasoning.update(icl_response) # reasoning first # Check if response is valid against json schema - correct_schema = json.loads(comparative_regression_json_schema(included_choices, target_kdma["factor"])) + + correct_schema = json.loads(comparative_regression_json_schema(schema_choices, target_kdma["factor"])) validate(instance=icl_response_with_reasoning, schema=correct_schema) # Get example prompt diff --git a/pyproject.toml b/pyproject.toml index 923fb43e..fe0238f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "pytest>=9.0.2", "rouge-score>=0.1.2", "scikit-learn>=1.7.2", - "swagger-client==0.6.1", + "swagger-client==0.6.2", "transformers>=4.56.0,<5", # Pinned to be compatible with vllm 0.19.0 "ubelt>=1.4.1", "setuptools-scm>=8.0,<9", diff --git a/tests/test_world_state_adm_component.py b/tests/test_world_state_adm_component.py new file mode 100644 index 00000000..6c675af4 --- /dev/null +++ b/tests/test_world_state_adm_component.py @@ -0,0 +1,229 @@ +import pytest +from copy import deepcopy + +from align_system.algorithms.world_state_adm_component import WorldStateTrackerADMComponent +from align_system.algorithms.pipeline_adm import PipelineADM +from align_system.algorithms.random_adm_component import RandomChoiceADMComponent + + +def test_world_state_tracker_initialization_and_update(): + component = WorldStateTrackerADMComponent() + + scene_1 = { + "unstructured": "Scene 1: Initial area.", + "elapsed_time": 0, + "meta_info": {"scene_id": "scene_1"}, + "supplies": [ + {"type": "Tourniquet", "quantity": 999}, + {"type": "Pressure bandage", "quantity": 999} + ], + "characters": [ + { + "id": "Patient 1", + "name": "Patient 1", + "unstructured": "Local civilian with moderate bleeding from their thigh", + "nearby": True, + "unseen": False, + "visited": False, + "tag": None, + "vitals": { + "avpu": "ALERT", + "breathing": "NORMAL", + "heart_rate": "FAST" + } + }, + { + "id": "Patient 2", + "name": "Patient 2", + "unstructured": "Military member of your unit with a broken left wrist, no bleeding", + "nearby": True, + "unseen": False, + "visited": False, + "tag": None, + "vitals": { + "avpu": "ALERT", + "breathing": "NORMAL", + "heart_rate": "FAST" + } + }, + { + "id": "Patient 6", + "name": "Patient 6", + "unstructured": "Civilian likely responsible for explosion", + "nearby": False, + "unseen": False, + "visited": False, + "tag": None + } + ] + } + + # First probe call + world_state, original_state, stale_info = component.run(scene_1, scenario_id="scenario_1") + + assert world_state["unstructured"] == "Scene 1: Initial area." + assert len(world_state["characters"]) == 3 + assert world_state["characters"][0]["vitals"]["avpu"] == "ALERT" + assert original_state == scene_1 + assert stale_info["stale_characters"] == {} + + # Second probe call: + # - Move away from Patient 1 (nearby becomes False, vitals omitted) + # - Patient 6 becomes nearby with vitals assessed + # - Tourniquet used (quantity decreases from 999 to 998) + scene_2 = { + "unstructured": "Scene 2: Moved to patient 6 area.", + "elapsed_time": 10, + "meta_info": {"scene_id": "scene_2"}, + "supplies": [ + {"type": "Tourniquet", "quantity": 998}, + {"type": "Pressure bandage", "quantity": 999} + ], + "characters": [ + { + "id": "Patient 1", + "name": "Patient 1", + "unstructured": "Patient 1 far away.", + "nearby": False, + "unseen": False, + "visited": True, # dynamic field updated + "tag": "RED" + # vitals omitted in scene_2 + }, + { + "id": "Patient 2", + "name": "Patient 2", + "unstructured": "Military member of your unit with a broken left wrist, no bleeding", + "nearby": True, + "unseen": False, + "visited": True, + "tag": None, + "vitals": { + "avpu": "ALERT", + "breathing": "NORMAL", + "heart_rate": "FAST" + } + }, + { + "id": "Patient 6", + "name": "Patient 6", + "unstructured": "Civilian likely responsible for the explosion with moderate bleeding from their torso; possible internal bleeding", + "nearby": True, + "unseen": False, + "visited": True, + "tag": None, + "vitals": { + "avpu": "VOICE", + "breathing": "FAST", + "heart_rate": "FAST" + } + } + ] + } + + world_state_2, original_state_2, stale_info_2 = component.run(scene_2, scenario_id="scenario_1") + + # Patient 1 should RETAIN detailed unstructured description AND vitals from scene 1 + p1 = next(c for c in world_state_2["characters"] if c["id"] == "Patient 1") + assert p1["unstructured"] == "Local civilian with moderate bleeding from their thigh" + assert p1["vitals"]["avpu"] == "ALERT" + assert p1["vitals"]["breathing"] == "NORMAL" + assert p1["visited"] is True + assert p1["tag"] == "RED" + assert p1["nearby"] is False + + # Patient 6 was nearby=False in scene 1, nearby=True in scene 2 -> Should update with new detailed description and vitals + p6 = next(c for c in world_state_2["characters"] if c["id"] == "Patient 6") + assert p6["unstructured"] == "Civilian likely responsible for the explosion with moderate bleeding from their torso; possible internal bleeding" + assert p6["vitals"]["avpu"] == "VOICE" + assert p6["nearby"] is True + + # Supplies quantity should be updated to 998 + t_supply = next(s for s in world_state_2["supplies"] if s["type"] == "Tourniquet") + assert t_supply["quantity"] == 998 + + # Stale info should track Patient 1 as having stale fields + assert "Patient 1" in stale_info_2["stale_characters"] + assert "unstructured" in stale_info_2["stale_characters"]["Patient 1"]["stale_fields"] + assert "vitals" in stale_info_2["stale_characters"]["Patient 1"]["stale_fields"] + + +def test_vitals_recursive_accumulation(): + component = WorldStateTrackerADMComponent() + + scene_1 = { + "characters": [ + { + "id": "P1", + "nearby": True, + "vitals": {"avpu": "ALERT"} + } + ] + } + + # Step 1: Initial vitals (avpu) + ws1, _, _ = component.run(scene_1, scenario_id="s1") + assert ws1["characters"][0]["vitals"] == {"avpu": "ALERT"} + + # Step 2: Assessment adds spo2 + scene_2 = { + "characters": [ + { + "id": "P1", + "nearby": True, + "vitals": {"spo2": "98%"} + } + ] + } + ws2, _, _ = component.run(scene_2, scenario_id="s1") + # Both avpu and spo2 should be present + assert ws2["characters"][0]["vitals"]["avpu"] == "ALERT" + assert ws2["characters"][0]["vitals"]["spo2"] == "98%" + + +def test_world_state_tracker_reset_history_and_scenario_change(): + component = WorldStateTrackerADMComponent() + + scene_s1 = { + "unstructured": "Scenario 1 Scene", + "characters": [{"id": "P1", "unstructured": "Details S1", "nearby": True}] + } + + component.run(scene_s1, scenario_id="scenario_1") + assert component.world_state is not None + + # Reset history + component.reset_history() + assert component.world_state is None + assert component.current_scenario_id is None + + # Auto reset on scenario ID change + component.run(scene_s1, scenario_id="scenario_1") + assert component.world_state["unstructured"] == "Scenario 1 Scene" + + scene_s2 = { + "unstructured": "Scenario 2 Scene", + "characters": [{"id": "P10", "unstructured": "Details S2", "nearby": True}] + } + + world_state_s2, _, _ = component.run(scene_s2, scenario_id="scenario_2") + assert world_state_s2["unstructured"] == "Scenario 2 Scene" + assert len(world_state_s2["characters"]) == 1 + assert world_state_s2["characters"][0]["id"] == "P10" + + +def test_pipeline_adm_reset_history_propagation(): + tracker = WorldStateTrackerADMComponent() + random_comp = RandomChoiceADMComponent() + + pipeline = PipelineADM(steps=[tracker, random_comp]) + + scene = {"unstructured": "Test", "characters": []} + tracker.run(scene, scenario_id="test_scen") + + assert tracker.world_state is not None + + # Call reset_history on PipelineADM + pipeline.reset_history() + + assert tracker.world_state is None