From 4646eb2f96eb2a7701f01240fa0ceb42035fa08e Mon Sep 17 00:00:00 2001 From: David Gardner Date: Mon, 23 Sep 2024 14:18:48 -0700 Subject: [PATCH] update tests to set the payload in the metadata --- tests/llm/nodes/test_extractor_node.py | 1 + tests/llm/nodes/test_manual_extractor_node.py | 1 + tests/llm/task_handlers/test_simple_task_handler.py | 1 + 3 files changed, 3 insertions(+) diff --git a/tests/llm/nodes/test_extractor_node.py b/tests/llm/nodes/test_extractor_node.py index 355f16aeb5..38673c95fc 100644 --- a/tests/llm/nodes/test_extractor_node.py +++ b/tests/llm/nodes/test_extractor_node.py @@ -39,6 +39,7 @@ def test_execute(): df = cudf.DataFrame({"insects": insects.copy(), "mammals": mammals.copy(), "reptiles": reptiles.copy()}) message = ControlMessage() message.payload(MessageMeta(df)) + message.set_metadata("llm_message_meta", message.payload()) task_dict = {"input_keys": ["mammals", "reptiles"]} node = ExtracterNode() diff --git a/tests/llm/nodes/test_manual_extractor_node.py b/tests/llm/nodes/test_manual_extractor_node.py index 143636999d..7c3ef3542e 100644 --- a/tests/llm/nodes/test_manual_extractor_node.py +++ b/tests/llm/nodes/test_manual_extractor_node.py @@ -48,6 +48,7 @@ def test_execute(): df = cudf.DataFrame({"insects": insects.copy(), "mammals": mammals.copy(), "reptiles": reptiles.copy()}) message = ControlMessage() message.payload(MessageMeta(df)) + message.set_metadata("llm_message_meta", message.payload()) task_dict = {"input_keys": ["insects"]} node = ManualExtracterNode(["mammals", "reptiles"]) diff --git a/tests/llm/task_handlers/test_simple_task_handler.py b/tests/llm/task_handlers/test_simple_task_handler.py index 8439d2df3d..641f65b9b4 100644 --- a/tests/llm/task_handlers/test_simple_task_handler.py +++ b/tests/llm/task_handlers/test_simple_task_handler.py @@ -46,6 +46,7 @@ def test_try_handle(dataset_cudf: DatasetManager): message = ControlMessage() message.payload(MessageMeta(df)) + message.set_metadata("llm_message_meta", message.payload()) task_handler = SimpleTaskHandler(['reptiles'])