🔧 refactor(base.py): add source_handle, target_handle, and target_param attributes to Edge class
The Edge class constructor now takes an additional parameter 'edge' which is a dictionary containing additional information about the edge. This information is used to set the source_handle, target_handle, and target_param attributes of the Edge class. This change allows for more flexibility and extensibility when working with edges in the graph.
🔧 refactor(base.py): add edge parameter to Edge class constructor to pass additional edge information
This commit is contained in:
parent
cfcdb67ed9
commit
766ef28a0c
2 changed files with 8 additions and 2 deletions
|
|
@ -6,9 +6,15 @@ if TYPE_CHECKING:
|
||||||
|
|
||||||
|
|
||||||
class Edge:
|
class Edge:
|
||||||
def __init__(self, source: "Vertex", target: "Vertex"):
|
def __init__(self, source: "Vertex", target: "Vertex", edge: dict):
|
||||||
self.source: "Vertex" = source
|
self.source: "Vertex" = source
|
||||||
self.target: "Vertex" = target
|
self.target: "Vertex" = target
|
||||||
|
self.source_handle = edge.get("sourceHandle", "")
|
||||||
|
self.target_handle = edge.get("targetHandle", "")
|
||||||
|
# 'BaseLoader;BaseOutputParser|documents|PromptTemplate-zmTlD'
|
||||||
|
# target_param is documents
|
||||||
|
self.target_param = self.target_handle.split("|")[1]
|
||||||
|
|
||||||
self.validate_edge()
|
self.validate_edge()
|
||||||
|
|
||||||
def validate_edge(self) -> None:
|
def validate_edge(self) -> None:
|
||||||
|
|
|
||||||
|
|
@ -179,7 +179,7 @@ class Graph:
|
||||||
raise ValueError(f"Source node {edge['source']} not found")
|
raise ValueError(f"Source node {edge['source']} not found")
|
||||||
if target is None:
|
if target is None:
|
||||||
raise ValueError(f"Target node {edge['target']} not found")
|
raise ValueError(f"Target node {edge['target']} not found")
|
||||||
edges.append(Edge(source, target))
|
edges.append(Edge(source, target, edge))
|
||||||
return edges
|
return edges
|
||||||
|
|
||||||
def _get_vertex_class(self, node_type: str, node_lc_type: str) -> Type[Vertex]:
|
def _get_vertex_class(self, node_type: str, node_lc_type: str) -> Type[Vertex]:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue