🔨 refactor(types.py): import flatten_list function from graph.utils module

This commit simply imports the flatten_list function from the graph.utils module to be used in the AgentVertex class. This improves the readability of the code and reduces the number of lines of code.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-02 14:14:47 -03:00
commit 534f47f153

View file

@ -1,15 +1,15 @@
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
from langflow.graph.vertex.base import Vertex from langflow.graph.vertex.base import Vertex
from langflow.graph.utils import extract_input_variables_from_prompt from langflow.graph.utils import extract_input_variables_from_prompt, flatten_list
class AgentVertex(Vertex): class AgentVertex(Vertex):
def __init__(self, data: Dict): def __init__(self, data: Dict):
super().__init__(data, base_type="agents") super().__init__(data, base_type="agents")
self.tools: List[Union[ToolNode, ToolkitNode]] = [] self.tools: List[Union[ToolVertex, ToolkitVertex]] = []
self.chains: List[ChainNode] = [] self.chains: List[ChainVertex] = []
def _set_tools_and_chains(self) -> None: def _set_tools_and_chains(self) -> None:
for edge in self.edges: for edge in self.edges:
@ -94,10 +94,10 @@ class ChainVertex(Vertex):
tools: Optional[Union[List[Vertex], List[ToolVertex]]] = None, tools: Optional[Union[List[Vertex], List[ToolVertex]]] = None,
) -> Any: ) -> Any:
if not self._built or force: if not self._built or force:
# Check if the chain requires a PromptNode # Check if the chain requires a PromptVertex
for key, value in self.params.items(): for key, value in self.params.items():
if isinstance(value, PromptVertex): if isinstance(value, PromptVertex):
# Build the PromptNode, passing the tools if available # Build the PromptVertex, passing the tools if available
self.params[key] = value.build(tools=tools, force=force) self.params[key] = value.build(tools=tools, force=force)
self._build() self._build()