Fix: fix query pipeline's arun_with_intermediates (#13002)

run-llama · Apr 22, 2024 · f1cefcb · f1cefcb
1 parent 521178f
commit f1cefcb
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/llama-index-core/llama_index/core/query_pipeline/query.py b/llama-index-core/llama_index/core/query_pipeline/query.py
@@ -778,11 +778,6 @@ async def _arun_multi(
                     [all_module_inputs[module_key] for module_key in popped_nodes],
                 )
 
-            if show_intermediates and module_key not in intermediate_outputs:
-                intermediate_outputs[module_key] = ComponentIntermediates(
-                    inputs=module_input, outputs=output_dict
-                )
-
             # create tasks from popped nodes
             tasks = []
             for module_key in popped_nodes:
@@ -801,6 +796,11 @@ async def _arun_multi(
                     queue, output_dict, module_key, all_module_inputs, result_outputs
                 )
 
+                if show_intermediates and module_key not in intermediate_outputs:
+                    intermediate_outputs[module_key] = ComponentIntermediates(
+                        inputs=all_module_inputs[module_key], outputs=output_dict
+                    )
+
         return result_outputs, intermediate_outputs
 
     def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]: