From 67764c01829fd5f6f661d7ab5e6889192c9cb67e Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Sun, 23 Jun 2024 15:55:30 -0300
Subject: [PATCH] Refactor constants.py to add Message type alias

---
 src/backend/base/langflow/field_typing/constants.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/backend/base/langflow/field_typing/constants.py b/src/backend/base/langflow/field_typing/constants.py
index 0156849cc..a8a0d60f3 100644
--- a/src/backend/base/langflow/field_typing/constants.py
+++ b/src/backend/base/langflow/field_typing/constants.py
@@ -1,4 +1,4 @@
-from typing import Callable, Dict, Text, TypeAlias, TypeVar, Union
+from typing import Callable, Dict, TypeAlias, TypeVar, Union
 
 from langchain.agents.agent import AgentExecutor
 from langchain.chains.base import Chain
@@ -16,6 +16,8 @@ from langchain_core.tools import Tool
 from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
 from langchain_text_splitters import TextSplitter
 
+from langflow.schema.message import Message
+
 NestedDict: TypeAlias = Dict[str, Union[str, Dict]]
 LanguageModel = TypeVar("LanguageModel", BaseLanguageModel, BaseLLM, BaseChatModel)
 Retriever = TypeVar(
@@ -23,6 +25,7 @@ Retriever = TypeVar(
     BaseRetriever,
     VectorStoreRetriever,
 )
+Text: TypeAlias = Message
 
 
 class Object:
@@ -62,6 +65,7 @@ CUSTOM_COMPONENT_SUPPORTED_TYPES = {
     **LANGCHAIN_BASE_TYPES,
     "NestedDict": NestedDict,
     "Data": Data,
+    "Message": Message,
     "Text": Text,
     "Object": Object,
     "Callable": Callable,