Return intermediate nodes output in pipelines

docs/_src/api/api/pipelines.md

@@ -121,6 +121,34 @@ Set the component for a node in the Pipeline.
 - `name`: The name of the node.
 - `component`: The component object to be set at the node.
 
+<a name="pipeline.Pipeline.run"></a>
+#### run
+
+```python
+ | run(query: Optional[str] = None, file_paths: Optional[List[str]] = None, labels: Optional[MultiLabel] = None, documents: Optional[List[Document]] = None, meta: Optional[dict] = None, params: Optional[dict] = None, debug: Optional[bool] = None, debug_logs: Optional[bool] = None)
+```
+
+Runs the pipeline, one node at a time.
+
+**Arguments**:
+
+- `query`: The search query (for query pipelines only)
+- `file_paths`: The files to index (for indexing pipelines only)
+- `labels`: 
+- `documents`: 
+- `meta`: 
+- `params`: Dictionary of parameters to be dispatched to the nodes.
+               If you want to pass a param to all nodes, you can just use: {"top_k":10}
+               If you want to pass it to targeted nodes, you can do:
+               {"Retriever": {"top_k": 10}, "Reader": {"top_k": 3, "debug": True}}
+- `debug`: Whether the pipeline should instruct nodes to collect debug information
+              about their execution. By default these include the input parameters
+              they received, the output they generated, and eventual logs (of any severity)
+              emitted. All debug information can then be found in the dict returned
+              by this method under the key "_debug"
+- `debug_logs`: Whether all the logs of the node should be printed in the console,
+                   regardless of their severity and of the existing logger's settings.
+
 <a name="pipeline.Pipeline.get_nodes_by_class"></a>
 #### get\_nodes\_by\_class
 

haystack/__init__.py

@@ -1,5 +1,14 @@
 import logging
 
+# Configure the root logger t0 DEBUG to allow the "debug" flag to receive the logs
+root_logger = logging.getLogger()
+root_logger.setLevel(logging.DEBUG)
+
+# Then reconfigure the StreamHandler not to display anything below WARNING as default
+stream_handler = logging.StreamHandler()
+stream_handler.setLevel(logging.WARNING)
+root_logger.addHandler(stream_handler)
+
 # Change log-levels before modules are loaded to avoid verbose log messages.
 logging.getLogger('haystack.modeling').setLevel(logging.WARNING)
 logging.getLogger('haystack.modeling.utils').setLevel(logging.INFO)

haystack/pipeline.py

@@ -260,7 +260,29 @@ def run(  # type: ignore
         documents: Optional[List[Document]] = None,
         meta: Optional[dict] = None,
         params: Optional[dict] = None,
+        debug: Optional[bool] = None,
+        debug_logs: Optional[bool] = None
     ):
+        """
+            Runs the pipeline, one node at a time.
+
+            :param query: The search query (for query pipelines only)
+            :param file_paths: The files to index (for indexing pipelines only)
+            :param labels: 
+            :param documents:
+            :param meta:
+            :param params: Dictionary of parameters to be dispatched to the nodes. 
+                           If you want to pass a param to all nodes, you can just use: {"top_k":10}
+                           If you want to pass it to targeted nodes, you can do:
+                           {"Retriever": {"top_k": 10}, "Reader": {"top_k": 3, "debug": True}}
+            :param debug: Whether the pipeline should instruct nodes to collect debug information
+                          about their execution. By default these include the input parameters
+                          they received, the output they generated, and eventual logs (of any severity)
+                          emitted. All debug information can then be found in the dict returned
+                          by this method under the key "_debug"
+            :param debug_logs: Whether all the logs of the node should be printed in the console,
+                               regardless of their severity and of the existing logger's settings.
+        """
         node_output = None
         queue = {
             self.root_node: {"root_node": self.root_node, "params": params}
@@ -281,6 +303,17 @@ def run(  # type: ignore
             node_id = list(queue.keys())[i]
             node_input = queue[node_id]
             node_input["node_id"] = node_id
+
+            # Apply debug attributes to the node input params
+            # NOTE: global debug attributes will override the value specified
+            # in each node's params dictionary.
+            if debug is not None:
+                if node_id not in node_input["params"].keys():
+                    node_input["params"][node_id] = {}
+                node_input["params"][node_id]["debug"] = debug
+                if debug_logs is not None:
+                    node_input["params"][node_id]["debug_logs"] = debug_logs
+
             predecessors = set(nx.ancestors(self.graph, node_id))
             if predecessors.isdisjoint(set(queue.keys())):  # only execute if predecessor nodes are executed
                 try:

haystack/schema.py

@@ -1,10 +1,15 @@
-from typing import Any, Optional, Dict, List
+from typing import Any, Optional, Dict, List, Callable, Tuple, Optional
+
 from uuid import uuid4
 from copy import deepcopy
 import mmh3
 import numpy as np
 from abc import abstractmethod
 import inspect
+import logging
+import io
+from functools import wraps
+
 
 class Document:
     def __init__(
@@ -249,6 +254,72 @@ def __str__(self):
         return str(self.to_dict())
 
 
+class InMemoryLogger(io.TextIOBase):
+    """
+    Implementation of a logger that keeps track
+    of the log lines in a list called `logs`,
+    from where they can be accessed freely.
+    """
+
+    def __init__(self, *args):
+        io.TextIOBase.__init__(self, *args)
+        self.logs = []
+
+    def write(self, x):
+        self.logs.append(x)
+
+
+def record_debug_logs(func: Callable, node_name: str, logs: bool) -> Callable:
+    """
+    Captures the debug logs of the wrapped function and 
+    saves them in the `_debug` key of the output dictionary. 
+    If `logs` is True, dumps the same logs to the console as well.
+
+    Used in `BaseComponent.__getattribute__()` to wrap `run()` functions.
+    This makes sure that every implementation of `run()` by a subclass will
+    be automagically decorated with this method when requested.
+
+    :param func: the function to decorate (must be an implementation of 
+                 `BaseComponent.run()`).
+    :param logs: whether the captured logs should also be displayed
+                 in the console during the execution of the pipeline.
+    """
+    @wraps(func)
+    def inner(*args, **kwargs) -> Tuple[Dict[str, Any], str]:
+
+        with InMemoryLogger() as logs_container:
+            logger = logging.getLogger()
+
+            # Adds a handler that stores the logs in a variable
+            handler = logging.StreamHandler(logs_container)
+            handler.setLevel(logger.level or logging.DEBUG)
+            logger.addHandler(handler)    
+
+            # Add a handler that prints log messages in the console
+            # to the specified level for the node
+            if logs:
+                handler_console = logging.StreamHandler()
+                handler_console.setLevel(logging.DEBUG)
+                formatter = logging.Formatter(f'[{node_name} logs] %(message)s')
+                handler_console.setFormatter(formatter)
+                logger.addHandler(handler_console)
+
+            output, stream = func(*args, **kwargs)
+
+            if not "_debug" in output.keys():
+                output["_debug"] = {}
+            output["_debug"]["logs"] = logs_container.logs
+
+            # Remove both handlers
+            logger.removeHandler(handler)
+            if logs:
+                logger.removeHandler(handler_console)
+
+            return output, stream
+
+    return inner
+
+
 class BaseComponent:
     """
     A base class for implementing nodes in a Pipeline.
@@ -266,6 +337,37 @@ def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         cls.subclasses[cls.__name__] = cls
 
+    def __getattribute__(self, name):
+        """
+        This modified `__getattribute__` method automagically decorates
+        every `BaseComponent.run()` implementation with the 
+        `record_debug_logs` decorator defined above. 
+
+        This decorator makes the function collect its debug logs into a 
+        `_debug` key of the output dictionary.
+
+        The logs collection is not always performed. Before applying the decorator,
+        it checks for an instance attribute called `debug` to know
+        whether it should or not. The decorator is applied if the attribute is 
+        defined and True.
+
+        In addition, the value of the instance attribute `debug_logs` is
+        passed to the decorator. If it's True, it will print the 
+        logs in the console as well.
+        """
+        if name == "run" and self.debug:
+            func = getattr(type(self), "run")
+            return record_debug_logs(func=func, node_name=self.__class__.__name__, logs=self.debug_logs).__get__(self)
+        return object.__getattribute__(self, name)
+
+    def __getattr__(self, name):
+        """
+        Ensures that `debug` and `debug_logs` are always defined.
+        """
+        if name in ["debug", "debug_logs"]:
+            return None
+        raise AttributeError(name)
+
     @classmethod
     def get_subclass(cls, component_type: str):
         if component_type not in cls.subclasses.keys():
@@ -317,7 +419,7 @@ def run(
         documents: Optional[List[Document]] = None,
         meta: Optional[dict] = None,
         params: Optional[dict] = None,
-    ):
+    ) -> Tuple[Dict, str]:
         """
         Method that will be executed when the node in the graph is called.
 
@@ -330,28 +432,37 @@ def run(
         """
         pass
 
-    def _dispatch_run(self, **kwargs):
+    def _dispatch_run(self, **kwargs) -> Tuple[Dict, str]:
         """
         The Pipelines call this method which in turn executes the run() method of Component.
 
         It takes care of the following:
           - inspect run() signature to validate if all necessary arguments are available
+          - pop `debug` and `debug_logs` and sets them on the instance to control debug output
           - call run() with the corresponding arguments and gather output
-          - collate _debug information if present
+          - collate `_debug` information if present
           - merge component output with the preceding output and pass it on to the subsequent Component in the Pipeline
         """
         arguments = deepcopy(kwargs)
         params = arguments.get("params") or {}
 
         run_signature_args = inspect.signature(self.run).parameters.keys()
 
-        run_params = {}
+        run_params: Dict[str, Any] = {}
         for key, value in params.items():
             if key == self.name:  # targeted params for this node
                 if isinstance(value, dict):
+
+                    # Extract debug attributes
+                    if "debug" in value.keys():
+                        self.debug = value.pop("debug")
+                    if "debug_logs" in value.keys():
+                        self.debug_logs = value.pop("debug_logs")
+
                     for _k, _v in value.items():
                         if _k not in run_signature_args:
                             raise Exception(f"Invalid parameter '{_k}' for the node '{self.name}'.")
+
                 run_params.update(**value)
             elif key in run_signature_args:  # global params
                 run_params[key] = value
@@ -363,9 +474,19 @@ def _dispatch_run(self, **kwargs):
 
         output, stream = self.run(**run_inputs, **run_params)
 
+        # Collect debug information
+        current_debug = output.get("_debug", {})
+        if self.debug:
+            current_debug["input"] = {**run_inputs, **run_params}
+            if self.debug:
+                current_debug["input"]["debug"] = self.debug
+            if self.debug_logs:
+                current_debug["input"]["debug_logs"] = self.debug_logs
+            filtered_output = {key: value for key, value in output.items() if key != "_debug"} # Exclude _debug to avoid recursion
+            current_debug["output"] = filtered_output        
+
         # append _debug information from nodes
         all_debug = arguments.get("_debug", {})
-        current_debug = output.get("_debug")
         if current_debug:
             all_debug[self.name] = current_debug
         if all_debug:

haystack/utils.py

@@ -1,3 +1,4 @@
+import io
 import json
 from collections import defaultdict
 from itertools import islice