From e8fded882e7294db31dab4fd29f8d182f8b329e6 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 25 Jan 2024 17:15:12 -0300 Subject: [PATCH] Add custom component and output parser creators to graph constants --- src/backend/langflow/graph/graph/constants.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/backend/langflow/graph/graph/constants.py b/src/backend/langflow/graph/graph/constants.py index b8617e56a..b18761de8 100644 --- a/src/backend/langflow/graph/graph/constants.py +++ b/src/backend/langflow/graph/graph/constants.py @@ -1,22 +1,21 @@ from langflow.graph.vertex import types from langflow.interface.agents.base import agent_creator from langflow.interface.chains.base import chain_creator +from langflow.interface.custom.base import custom_component_creator from langflow.interface.document_loaders.base import documentloader_creator from langflow.interface.embeddings.base import embedding_creator from langflow.interface.llms.base import llm_creator from langflow.interface.memories.base import memory_creator +from langflow.interface.output_parsers.base import output_parser_creator from langflow.interface.prompts.base import prompt_creator +from langflow.interface.retrievers.base import retriever_creator from langflow.interface.text_splitters.base import textsplitter_creator from langflow.interface.toolkits.base import toolkits_creator from langflow.interface.tools.base import tool_creator from langflow.interface.vector_store.base import vectorstore_creator from langflow.interface.wrappers.base import wrapper_creator -from langflow.interface.output_parsers.base import output_parser_creator -from langflow.interface.retrievers.base import retriever_creator -from langflow.interface.custom.base import custom_component_creator from langflow.utils.lazy_load import LazyLoadDictBase - chat_components = ["ChatInput", "ChatOutput", "TextInput", "SessionID"] @@ -50,13 +49,13 @@ class VertexTypesDict(LazyLoadDictBase): **{t: types.DocumentLoaderVertex for t in documentloader_creator.to_list()}, **{t: types.TextSplitterVertex for t in textsplitter_creator.to_list()}, **{t: types.OutputParserVertex for t in output_parser_creator.to_list()}, - **{ - t: types.CustomComponentVertex - for t in custom_component_creator.to_list() - }, + **{t: types.CustomComponentVertex for t in custom_component_creator.to_list()}, **{t: types.RetrieverVertex for t in retriever_creator.to_list()}, **{t: types.ChatVertex for t in chat_components}, } + def get_custom_component_vertex_type(self): + return types.CustomComponentVertex + lazy_load_vertex_dict = VertexTypesDict()