From 566c48166ed59ba2f43a0dca38e11484046f7f30 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 31 Jul 2023 10:35:48 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=80=20refactor(custom=5Fcomponent.py):?= =?UTF-8?q?=20modify=20get=5Fflow=20method=20to=20accept=20keyword=20argum?= =?UTF-8?q?ents=20for=20flow=5Fname=20and=20flow=5Fid=20to=20improve=20fle?= =?UTF-8?q?xibility=20and=20readability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langflow/interface/custom/custom_component.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index bb3bafb08..e1dcd83ab 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -158,14 +158,12 @@ class CustomComponent(Component, extra=Extra.allow): flows = session.query(Flow).all() return flows - def get_flow(self, flow_name: str) -> Flow: + def get_flow(self, *, flow_name: str = None, flow_id: str = None) -> Flow: with session_getter() as session: - flow = session.query(Flow).filter(Flow.name == flow_name).first() - return flow - - def get_flow_by_id(self, flow_id: str) -> Flow: - with session_getter() as session: - flow = session.query(Flow).filter(Flow.id == flow_id).first() + if flow_id: + flow = session.query(Flow).get(flow_id) + elif flow_name: + flow = session.query(Flow).filter(Flow.name == flow_name).first() return flow def build(self):