From f1de9a554da3f2f46be85ae55fb0cde08be78e30 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 29 Jun 2023 11:15:51 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=80=20refactor(constants.py):=20import?=
 =?UTF-8?q?=20vertex=20types=20from=20langflow.graph.vertex.types=20as=20t?=
 =?UTF-8?q?ypes=20for=20improved=20readability=20and=20maintainability?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The import statements for vertex types in the constants.py file have been refactored to import from langflow.graph.vertex.types as types. This change improves the readability and maintainability of the code by making it clear where the vertex types are coming from.
---
 src/backend/langflow/graph/graph/constants.py | 42 +++++++------------
 1 file changed, 15 insertions(+), 27 deletions(-)

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()},
 }