Refactor Graph class to improve readability and add docstrings

src/backend/langflow/graph/graph/base.py

@@ -35,6 +35,14 @@ class Graph:
         edges: List[Dict[str, str]],
         flow_id: Optional[str] = None,
     ) -> None:
+        """
+        Initializes a new instance of the Graph class.
+
+        Args:
+            nodes (List[Dict]): A list of dictionaries representing the vertices of the graph.
+            edges (List[Dict[str, str]]): A list of dictionaries representing the edges of the graph.
+            flow_id (Optional[str], optional): The ID of the flow. Defaults to None.
+        """
         self._vertices = nodes
         self._edges = edges
         self.raw_graph_data = {"nodes": nodes, "edges": edges}
@@ -71,11 +79,26 @@ class Graph:
         self.state_manager = GraphStateManager()
 
     def get_state(self, name: str) -> Optional[Record]:
-        """Returns the state of the graph."""
+        """
+        Returns the state of the graph with the given name.
+
+        Args:
+            name (str): The name of the state.
+
+        Returns:
+            Optional[Record]: The state record, or None if the state does not exist.
+        """
         return self.state_manager.get_state(name, run_id=self._run_id)
 
     def update_state(self, name: str, record: Union[str, Record], caller: Optional[str] = None) -> None:
-        """Updates the state of the graph."""
+        """
+        Updates the state of the graph with the given name.
+
+        Args:
+            name (str): The name of the state.
+            record (Union[str, Record]): The new state record.
+            caller (Optional[str], optional): The ID of the vertex that is updating the state. Defaults to None.
+        """
         if caller:
             # If there is a caller which is a vertex_id, I want to activate
             # all StateVertex in self.vertices that are not the caller
@@ -86,6 +109,13 @@ class Graph:
         self.state_manager.update_state(name, record, run_id=self._run_id)
 
     def activate_state_vertices(self, name: str, caller: str):
+        """
+        Activates the state vertices in the graph with the given name and caller.
+
+        Args:
+            name (str): The name of the state.
+            caller (str): The ID of the vertex that is updating the state.
+        """
         vertices_ids = []
         for vertex_id in self._is_state_vertices:
             if vertex_id == caller:
@@ -104,10 +134,20 @@ class Graph:
         self.vertices_to_run.update(vertices_ids)
 
     def reset_activated_vertices(self):
+        """
+        Resets the activated vertices in the graph.
+        """
         self.activated_vertices = []
 
     def append_state(self, name: str, record: Union[str, Record], caller: Optional[str] = None) -> None:
-        """Appends the state of the graph."""
+        """
+        Appends the state of the graph with the given name.
+
+        Args:
+            name (str): The name of the state.
+            record (Union[str, Record]): The state record to append.
+            caller (Optional[str], optional): The ID of the vertex that is updating the state. Defaults to None.
+        """
         if caller:
             self.activate_state_vertices(name, caller)
 
@@ -115,17 +155,38 @@ class Graph:
 
     @property
     def run_id(self):
+        """
+        The ID of the current run.
+
+        Returns:
+            str: The run ID.
+
+        Raises:
+            ValueError: If the run ID is not set.
+        """
         if not self._run_id:
             raise ValueError("Run ID not set")
         return self._run_id
 
     def set_run_id(self, run_id: str):
+        """
+        Sets the ID of the current run.
+
+        Args:
+            run_id (str): The run ID.
+        """
         for vertex in self.vertices:
             self.state_manager.subscribe(run_id, vertex.update_graph_state)
         self._run_id = run_id
 
     @property
     def sorted_vertices_layers(self) -> List[List[str]]:
+        """
+        The sorted layers of vertices in the graph.
+
+        Returns:
+            List[List[str]]: The sorted layers of vertices.
+        """
         if not self._sorted_vertices_layers:
             self.sort_vertices()
         return self._sorted_vertices_layers
@@ -148,7 +209,19 @@ class Graph:
         stream: bool,
         session_id: str,
     ) -> List[Optional["ResultData"]]:
-        """Runs the graph with the given inputs."""
+        """
+        Runs the graph with the given inputs.
+
+        Args:
+            inputs (Dict[str, str]): The input values for the graph.
+            input_components (list[str]): The components to run for the inputs.
+            outputs (list[str]): The outputs to retrieve from the graph.
+            stream (bool): Whether to stream the results or not.
+            session_id (str): The session ID for the graph.
+
+        Returns:
+            List[Optional["ResultData"]]: The outputs of the graph.
+        """
         for vertex_id in self._is_input_vertices:
             vertex = self.get_vertex(vertex_id)
             if input_components and (vertex_id not in input_components or vertex.display_name not in input_components):
@@ -190,7 +263,19 @@ class Graph:
         session_id: Optional[str] = None,
         stream: bool = False,
     ) -> List[RunOutputs]:
-        """Runs the graph with the given inputs."""
+        """
+        Runs the graph with the given inputs.
+
+        Args:
+            inputs (list[Dict[str, str]]): The input values for the graph.
+            inputs_components (Optional[list[list[str]]], optional): The components to run for the inputs. Defaults to None.
+            outputs (Optional[list[str]], optional): The outputs to retrieve from the graph. Defaults to None.
+            session_id (Optional[str], optional): The session ID for the graph. Defaults to None.
+            stream (bool, optional): Whether to stream the results or not. Defaults to False.
+
+        Returns:
+            List[RunOutputs]: The outputs of the graph.
+        """
         # inputs is {"message": "Hello, world!"}
         # we need to go through self.inputs and update the self._raw_params
         # of the vertices that are inputs
@@ -218,16 +303,23 @@ class Graph:
             vertex_outputs.append(run_output_object)
         return vertex_outputs
 
-    # vertices_layers is a list of lists ordered by the order the vertices
-    # should be built.
-    # We need to create a new method that will take the vertices_layers
-    # and return the next vertex to be built.
     def next_vertex_to_build(self):
-        """Returns the next vertex to be built."""
+        """
+        Returns the next vertex to be built.
+
+        Yields:
+            str: The ID of the next vertex to be built.
+        """
         yield from chain.from_iterable(self.vertices_layers)
 
     @property
     def metadata(self):
+        """
+        The metadata of the graph.
+
+        Returns:
+            dict: The metadata of the graph.
+        """
         return {
             "runs": self._runs,
             "updates": self._updates,
@@ -235,12 +327,19 @@ class Graph:
         }
 
     def build_graph_maps(self):
+        """
+        Builds the adjacency maps for the graph.
+        """
         self.predecessor_map, self.successor_map = self.build_adjacency_maps()
 
         self.in_degree_map = self.build_in_degree()
         self.parent_child_map = self.build_parent_child_map()
 
     def reset_inactivated_vertices(self):
+        """
+        Resets the inactivated vertices in the graph.
+        """
+        self.inactivated_vertices = []
         self.inactivated_vertices = set()
 
     def mark_all_vertices(self, state: str):