Add is_state attribute to Vertex class and StateVertex subclass

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-02 01:18:58 -03:00
commit 2b215e66f4
2 changed files with 3 additions and 2 deletions

View file

@ -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
)

View file

@ -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: