jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

🐛 fix(custom_component.py): fix session handling in load_flow method to ensure proper context management and avoid potential resource leaks

Browse source 
✨ feat(custom_component.py): add custom_repr method to CustomComponent class to provide a custom representation value for the component
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-07-21 09:25:55 -03:00
commit 4307c24c6d
1 changed files with 9 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

13
src/backend/langflow/interface/custom/custom_component.py
View file

@ -6,21 +6,26 @@ from langflow.interface.custom.component import Component
from langflow.utils import validate from langflow.utils import validate
from uuid import UUID from uuid import UUID
from langflow.database.base import get_session from langflow.database.base import session_getter
from langflow.database.models.flow import Flow from langflow.database.models.flow import Flow
from pydantic import Extra
class CustomComponent(Component): class CustomComponent(Component, extra=Extra.allow):
code: Optional[str] code: Optional[str]
field_config: dict = {} field_config: dict = {}
code_class_base_inheritance = "CustomComponent" code_class_base_inheritance = "CustomComponent"
function_entrypoint_name = "build" function_entrypoint_name = "build"
function: Optional[Callable] = None function: Optional[Callable] = None
return_type_valid_list = list(LANGCHAIN_BASE_TYPES.keys()) return_type_valid_list = list(LANGCHAIN_BASE_TYPES.keys())
repr_value: Optional[str] = ""
def __init__(self, **data): def __init__(self, **data):
super().__init__(**data) super().__init__(**data)
def custom_repr(self):
return self.repr_value
def _class_template_validation(self, code: str) -> bool: def _class_template_validation(self, code: str) -> bool:
if not code: if not code:
raise HTTPException( raise HTTPException(
@ -133,8 +138,8 @@ class CustomComponent(Component):
def load_flow(self, flow_id: UUID = None): def load_flow(self, flow_id: UUID = None):
from langflow.processing.process import build_sorted_vertices_with_caching from langflow.processing.process import build_sorted_vertices_with_caching
session = next(get_session()) with session_getter() as session:
data_graph = flow.data if (flow := session.get(Flow, flow_id)) else None data_graph = flow.data if (flow := session.get(Flow, flow_id)) else None
if not data_graph: if not data_graph:
raise ValueError(f"Flow {flow_id} not found") raise ValueError(f"Flow {flow_id} not found")
return build_sorted_vertices_with_caching(data_graph) return build_sorted_vertices_with_caching(data_graph)