diff --git a/src/backend/langflow/graph/graph/constants.py b/src/backend/langflow/graph/graph/constants.py index 3398f253f..9474e311e 100644 --- a/src/backend/langflow/graph/graph/constants.py +++ b/src/backend/langflow/graph/graph/constants.py @@ -1,18 +1,5 @@ from langflow.graph.vertex.base import Vertex -from langflow.graph.vertex.types import ( - AgentVertex, - ChainVertex, - DocumentLoaderVertex, - EmbeddingVertex, - LLMVertex, - MemoryVertex, - PromptVertex, - TextSplitterVertex, - ToolVertex, - ToolkitVertex, - VectorStoreVertex, - WrapperVertex, -) +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.document_loaders.base import documentloader_creator @@ -25,22 +12,23 @@ 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.retrievers.base import retriever_creator from typing import Dict, Type VERTEX_TYPE_MAP: Dict[str, Type[Vertex]] = { - **{t: PromptVertex for t in prompt_creator.to_list()}, - **{t: AgentVertex for t in agent_creator.to_list()}, - **{t: ChainVertex for t in chain_creator.to_list()}, - **{t: ToolVertex for t in tool_creator.to_list()}, - **{t: ToolkitVertex for t in toolkits_creator.to_list()}, - **{t: WrapperVertex for t in wrapper_creator.to_list()}, - **{t: LLMVertex for t in llm_creator.to_list()}, - **{t: MemoryVertex for t in memory_creator.to_list()}, - **{t: EmbeddingVertex for t in embedding_creator.to_list()}, - **{t: VectorStoreVertex for t in vectorstore_creator.to_list()}, - **{t: DocumentLoaderVertex for t in documentloader_creator.to_list()}, - **{t: TextSplitterVertex for t in textsplitter_creator.to_list()}, + **{t: types.PromptVertex for t in prompt_creator.to_list()}, + **{t: types.AgentVertex for t in agent_creator.to_list()}, + **{t: types.ChainVertex for t in chain_creator.to_list()}, + **{t: types.ToolVertex for t in tool_creator.to_list()}, + **{t: types.ToolkitVertex for t in toolkits_creator.to_list()}, + **{t: types.WrapperVertex for t in wrapper_creator.to_list()}, + **{t: types.LLMVertex for t in llm_creator.to_list()}, + **{t: types.MemoryVertex for t in memory_creator.to_list()}, + **{t: types.EmbeddingVertex for t in embedding_creator.to_list()}, + **{t: types.VectorStoreVertex for t in vectorstore_creator.to_list()}, + **{t: types.DocumentLoaderVertex for t in documentloader_creator.to_list()}, + **{t: types.TextSplitterVertex for t in textsplitter_creator.to_list()}, + **{t: types.RetrieverVertex for t in retriever_creator.to_list()}, }