diff --git a/biochatter/prompts.py b/biochatter/prompts.py index 8c62ec69..b7bd1132 100644 --- a/biochatter/prompts.py +++ b/biochatter/prompts.py @@ -372,16 +372,14 @@ def _select_relationships(self) -> bool: if self.selected_relationship_labels: for key, value in self.selected_relationship_labels.items(): - if ( - value["source"] - and value["source"] not in self.selected_entities - ): - self.selected_entities.append(value["source"]) - if ( - value["target"] - and value["target"] not in self.selected_entities - ): - self.selected_entities.append(value["target"]) + sources = ensure_iterable(value["source"]) + targets = ensure_iterable(value["target"]) + for source in sources: + if source not in self.selected_entities: + self.selected_entities.append(source) + for target in targets: + if target not in self.selected_entities: + self.selected_entities.append(target) return bool(result)