Fix unit tests magpie

Author: gabrielmbmb

Diff for: tests/unit/steps/tasks/magpie/test_base.py

@@ -30,14 +30,14 @@ def test_raise_value_error_llm_no_magpie_mixin(self) -> None:
     def test_outputs(self) -> None:
         task = Magpie(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"))
 
-        assert task.outputs == ["conversation"]
+        assert task.outputs == ["conversation", "model_name"]
 
         task = Magpie(
             llm=DummyMagpieLLM(magpie_pre_query_template="llama3"),
             only_instruction=True,
         )
 
-        assert task.outputs == ["instruction"]
+        assert task.outputs == ["instruction", "model_name"]
 
     def test_process(self) -> None:
         task = Magpie(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"), n_turns=1)
@@ -50,18 +50,21 @@ def test_process(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
         ]
 
@@ -79,6 +82,7 @@ def test_process_with_n_turns(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
@@ -88,6 +92,7 @@ def test_process_with_n_turns(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
                 "conversation": [
@@ -97,6 +102,7 @@ def test_process_with_n_turns(self) -> None:
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
         ]
 
@@ -115,31 +121,37 @@ def test_process_with_system_prompt_per_row(self) -> None:
             )
         ) == [
             {
+                "system_prompt": "You're a math expert assistant.",
                 "conversation": [
                     {"role": "system", "content": "You're a math expert assistant."},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
+                "system_prompt": "You're a florist expert assistant.",
                 "conversation": [
                     {"role": "system", "content": "You're a florist expert assistant."},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
             {
+                "system_prompt": "You're a plumber expert assistant.",
                 "conversation": [
                     {"role": "system", "content": "You're a plumber expert assistant."},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                     {"role": "user", "content": "Hello Magpie"},
                     {"role": "assistant", "content": "Hello Magpie"},
                 ],
+                "model_name": "test",
             },
         ]
 
@@ -152,9 +164,18 @@ def test_process_only_instruction(self) -> None:
         task.load()
 
         assert next(task.process(inputs=[{}, {}, {}])) == [
-            {"instruction": "Hello Magpie"},
-            {"instruction": "Hello Magpie"},
-            {"instruction": "Hello Magpie"},
+            {
+                "instruction": "Hello Magpie",
+                "model_name": "test",
+            },
+            {
+                "instruction": "Hello Magpie",
+                "model_name": "test",
+            },
+            {
+                "instruction": "Hello Magpie",
+                "model_name": "test",
+            },
         ]
 
     def test_serialization(self) -> None:

Diff for: tests/unit/steps/tasks/magpie/test_generator.py

@@ -30,14 +30,14 @@ def test_raise_value_error_llm_no_magpie_mixin(self) -> None:
     def test_outputs(self) -> None:
         task = MagpieGenerator(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"))
 
-        assert task.outputs == ["conversation"]
+        assert task.outputs == ["conversation", "model_name"]
 
         task = MagpieGenerator(
             llm=DummyMagpieLLM(magpie_pre_query_template="llama3"),
             only_instruction=True,
         )
 
-        assert task.outputs == ["instruction"]
+        assert task.outputs == ["instruction", "model_name"]
 
     def test_serialization(self) -> None:
         task = MagpieGenerator(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"))