From 2b215e66f4ccfb340125af887b0bcff0499ff86e Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 2 Mar 2024 01:18:58 -0300 Subject: [PATCH] Add is_state attribute to Vertex class and StateVertex subclass --- src/backend/langflow/graph/vertex/base.py | 1 + src/backend/langflow/graph/vertex/types.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py index 2f11e0ed3..097b99fbc 100644 --- a/src/backend/langflow/graph/vertex/base.py +++ b/src/backend/langflow/graph/vertex/base.py @@ -57,6 +57,7 @@ class Vertex: self.will_stream = False self.updated_raw_params = False self.id: str = data["id"] + self.is_state = False self.is_input = any( input_component_name in self.id for input_component_name in INPUT_COMPONENTS ) diff --git a/src/backend/langflow/graph/vertex/types.py b/src/backend/langflow/graph/vertex/types.py index ba1cd2998..516b06dee 100644 --- a/src/backend/langflow/graph/vertex/types.py +++ b/src/backend/langflow/graph/vertex/types.py @@ -1,7 +1,6 @@ import ast import json -from typing import (AsyncIterator, Callable, Dict, Iterator, List, Optional, - Union) +from typing import AsyncIterator, Callable, Dict, Iterator, List, Optional, Union import yaml from langchain_core.messages import AIMessage @@ -487,6 +486,7 @@ class StateVertex(Vertex): def __init__(self, data: Dict, graph): super().__init__(data, graph=graph, base_type="custom_components") self.steps = [self._build] + self.is_state = True def dict_to_codeblock(d: dict) -> str: