Add end_with_user and include_system_prompt flags to Magpie tasks and handle Nones.

src/distilabel/steps/tasks/magpie/base.py

@@ -50,6 +50,10 @@ class MagpieBase(RuntimeParametersMixin):
         default=1,
         description="The number of turns to generate for the conversation.",
     )
+    end_with_user: RuntimeParameter[bool] = Field(
+        default=False,
+        description="Whether the conversation should end with a user message.",
+    )
     only_instruction: RuntimeParameter[bool] = Field(
         default=False,
         description="Whether to generate only the instruction. If this argument"
@@ -125,7 +129,7 @@ def _generate_multi_turn_conversation(
     ) -> List[Dict[str, Any]]:
         conversations = self._prepare_inputs_for_instruction_generation(inputs)
 
-        for _ in range(self.n_turns):  # type: ignore
+        for i in range(self.n_turns):  # type: ignore
             # Generate instruction or user message
             outputs = self.llm.generate(
                 inputs=conversations,
@@ -139,6 +143,9 @@ def _generate_multi_turn_conversation(
                 conversations=conversations,  # type: ignore
             )
 
+            if i == self.n_turns - 1 and self.end_with_user:  # type: ignore
+                break
+
             # TODO: handle potential previous `None`s
 
             # Generate response

tests/unit/steps/tasks/magpie/test_base.py

@@ -106,6 +106,45 @@ def test_process_with_n_turns(self) -> None:
             },
         ]
 
+    def test_process_with_end_with_user(self) -> None:
+        task = Magpie(
+            llm=DummyMagpieLLM(magpie_pre_query_template="llama3"),
+            n_turns=2,
+            end_with_user=True,
+        )
+
+        task.load()
+
+        assert next(task.process(inputs=[{}, {}, {}])) == [
+            {
+                "conversation": [
+                    {"role": "system", "content": MAGPIE_MULTI_TURN_SYSTEM_PROMPT},
+                    {"role": "user", "content": "Hello Magpie"},
+                    {"role": "assistant", "content": "Hello Magpie"},
+                    {"role": "user", "content": "Hello Magpie"},
+                ],
+                "model_name": "test",
+            },
+            {
+                "conversation": [
+                    {"role": "system", "content": MAGPIE_MULTI_TURN_SYSTEM_PROMPT},
+                    {"role": "user", "content": "Hello Magpie"},
+                    {"role": "assistant", "content": "Hello Magpie"},
+                    {"role": "user", "content": "Hello Magpie"},
+                ],
+                "model_name": "test",
+            },
+            {
+                "conversation": [
+                    {"role": "system", "content": MAGPIE_MULTI_TURN_SYSTEM_PROMPT},
+                    {"role": "user", "content": "Hello Magpie"},
+                    {"role": "assistant", "content": "Hello Magpie"},
+                    {"role": "user", "content": "Hello Magpie"},
+                ],
+                "model_name": "test",
+            },
+        ]
+
     def test_process_with_system_prompt_per_row(self) -> None:
         task = Magpie(llm=DummyMagpieLLM(magpie_pre_query_template="llama3"), n_turns=2)
 
@@ -195,6 +234,7 @@ def test_serialization(self) -> None:
                 },
             },
             "n_turns": 1,
+            "end_with_user": False,
             "only_instruction": True,
             "system_prompt": None,
             "name": "magpie_0",
@@ -227,6 +267,11 @@ def test_serialization(self) -> None:
                     "optional": True,
                     "description": "The number of turns to generate for the conversation.",
                 },
+                {
+                    "name": "end_with_user",
+                    "optional": True,
+                    "description": "Whether the conversation should end with a user message.",
+                },
                 {
                     "name": "only_instruction",
                     "optional": True,

tests/unit/steps/tasks/magpie/test_generator.py

@@ -53,6 +53,7 @@ def test_serialization(self) -> None:
                 },
             },
             "n_turns": 1,
+            "end_with_user": False,
             "only_instruction": False,
             "system_prompt": None,
             "name": "magpie_generator_0",
@@ -86,6 +87,11 @@ def test_serialization(self) -> None:
                     "optional": True,
                     "description": "The number of turns to generate for the conversation.",
                 },
+                {
+                    "name": "end_with_user",
+                    "optional": True,
+                    "description": "Whether the conversation should end with a user message.",
+                },
                 {
                     "name": "only_instruction",
                     "optional": True,