allow non leaf outputs in supercomponents

Author: mdrazak2001

haystack/core/super_component/super_component.py

@@ -129,7 +129,7 @@ def __init__(
         self._original_input_mapping = input_mapping
 
         # Set output types based on pipeline and mapping
-        pipeline_outputs = self.pipeline.outputs()
+        pipeline_outputs = self.pipeline.outputs(include_components_with_connected_outputs=True)
         resolved_output_mapping = (
             output_mapping if output_mapping is not None else self._create_output_mapping(pipeline_outputs)
         )
@@ -189,9 +189,14 @@ def run(self, **kwargs: Any) -> Dict[str, Any]:
         """
         filtered_inputs = {param: value for param, value in kwargs.items() if value != _delegate_default}
         pipeline_inputs = self._map_explicit_inputs(input_mapping=self.input_mapping, inputs=filtered_inputs)
-        pipeline_outputs = self.pipeline.run(data=pipeline_inputs)
+        include_outputs_from = self._get_include_outputs_from()
+        pipeline_outputs = self.pipeline.run(data=pipeline_inputs, include_outputs_from=include_outputs_from)
         return self._map_explicit_outputs(pipeline_outputs, self.output_mapping)
 
+    def _get_include_outputs_from(self) -> set[str]:
+        # Collecting the component names from output_mapping
+        return {self._split_component_path(path)[0] for path in self.output_mapping.keys()}
+
     async def run_async(self, **kwargs: Any) -> Dict[str, Any]:
         """
         Runs the wrapped pipeline with the provided inputs async.

test/core/super_component/test_super_component.py

@@ -278,3 +278,13 @@ def from_dict(cls, data):
 
         assert custom_serialized["type"] == "test_super_component.CustomSuperComponent"
         assert custom_super_component._to_super_component_dict() == serialized
+
+    def test_super_component_non_leaf_output(self, rag_pipeline):
+        # 'retriever' is not a leaf, but should now be allowed
+        output_mapping = {"retriever.documents": "retrieved_docs", "answer_builder.answers": "final_answers"}
+        wrapper = SuperComponent(pipeline=rag_pipeline, output_mapping=output_mapping)
+        wrapper.warm_up()
+        result = wrapper.run(query="What is the capital of France?")
+        assert "final_answers" in result  # leaf output
+        assert "retrieved_docs" in result  # non-leaf output
+        assert isinstance(result["retrieved_docs"][0], Document)