From 766ef28a0c1d278f554f41e06c65810f0371d996 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 29 Jun 2023 22:06:57 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20refactor(base.py):=20add=20sourc?= =?UTF-8?q?e=5Fhandle,=20target=5Fhandle,=20and=20target=5Fparam=20attribu?= =?UTF-8?q?tes=20to=20Edge=20class=20The=20Edge=20class=20constructor=20no?= =?UTF-8?q?w=20takes=20an=20additional=20parameter=20'edge'=20which=20is?= =?UTF-8?q?=20a=20dictionary=20containing=20additional=20information=20abo?= =?UTF-8?q?ut=20the=20edge.=20This=20information=20is=20used=20to=20set=20?= =?UTF-8?q?the=20source=5Fhandle,=20target=5Fhandle,=20and=20target=5Fpara?= =?UTF-8?q?m=20attributes=20of=20the=20Edge=20class.=20This=20change=20all?= =?UTF-8?q?ows=20for=20more=20flexibility=20and=20extensibility=20when=20w?= =?UTF-8?q?orking=20with=20edges=20in=20the=20graph.=20=F0=9F=94=A7=20refa?= =?UTF-8?q?ctor(base.py):=20add=20edge=20parameter=20to=20Edge=20class=20c?= =?UTF-8?q?onstructor=20to=20pass=20additional=20edge=20information?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/graph/edge/base.py | 8 +++++++- src/backend/langflow/graph/graph/base.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/graph/edge/base.py b/src/backend/langflow/graph/edge/base.py index 08f084a5c..341c3c78f 100644 --- a/src/backend/langflow/graph/edge/base.py +++ b/src/backend/langflow/graph/edge/base.py @@ -6,9 +6,15 @@ if TYPE_CHECKING: class Edge: - def __init__(self, source: "Vertex", target: "Vertex"): + def __init__(self, source: "Vertex", target: "Vertex", edge: dict): self.source: "Vertex" = source 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() def validate_edge(self) -> None: diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index 46425ddf6..65fff3239 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -179,7 +179,7 @@ class Graph: raise ValueError(f"Source node {edge['source']} not found") if target is None: raise ValueError(f"Target node {edge['target']} not found") - edges.append(Edge(source, target)) + edges.append(Edge(source, target, edge)) return edges def _get_vertex_class(self, node_type: str, node_lc_type: str) -> Type[Vertex]: