Coverage for haystack/core/pipeline/descriptions.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5 

6import networkx 

7 

8from haystack.core.component.types import InputSocket, OutputSocket 

9 

10 

11def find_pipeline_inputs( 

12 graph: networkx.MultiDiGraph, include_connected_sockets: bool = False 

13) -> dict[str, list[InputSocket]]: 

14 """ 

15 Collect components that have disconnected/connected input sockets. 

16 

17 Note that this method returns *ALL* disconnected input sockets, including all such sockets with default values. 

18 It also includes variadic input sockets, even if they are currently connected, as they can accept additional 

19 inputs from outside the pipeline. 

20 

21 :param graph: The pipeline graph to analyze. 

22 :param include_connected_sockets: If True, also include input sockets that are already connected. 

23 This can be useful for understanding the full input requirements of the pipeline, including inputs 

24 that are currently satisfied by connections within the pipeline. If False, only include input sockets that 

25 are not connected to any output socket, which represent the external inputs that can be provided when running 

26 the pipeline. 

27 """ 

28 return { 

29 name: [ 

30 socket 

31 for socket in data.get("input_sockets", {}).values() 

32 if socket.is_variadic or (include_connected_sockets or not socket.senders) 

33 ] 

34 for name, data in graph.nodes(data=True) 

35 } 

36 

37 

38def find_pipeline_outputs( 

39 graph: networkx.MultiDiGraph, include_connected_sockets: bool = False 

40) -> dict[str, list[OutputSocket]]: 

41 """ 

42 Collect components that have disconnected/connected output sockets. They define the pipeline output. 

43 """ 

44 return { 

45 name: [ 

46 socket 

47 for socket in data.get("output_sockets", {}).values() 

48 if (include_connected_sockets or not socket.receivers) 

49 ] 

50 for name, data in graph.nodes(data=True) 

51 }