From df3c275db89762582465af2bdb5330fa547810dd Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Thu, 20 Jun 2024 19:05:16 -0300
Subject: [PATCH] refactor: Update langflow schema data module to use
 BaseMessage instead of HumanMessage or SystemMessage

---
 src/backend/base/langflow/schema/data.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/backend/base/langflow/schema/data.py b/src/backend/base/langflow/schema/data.py
index 919b165d9..122f01507 100644
--- a/src/backend/base/langflow/schema/data.py
+++ b/src/backend/base/langflow/schema/data.py
@@ -3,7 +3,7 @@ import json
 from typing import Optional, cast
 
 from langchain_core.documents import Document
-from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
 from langchain_core.prompt_values import ImagePromptValue
 from langchain_core.prompts.image import ImagePromptTemplate
 from pydantic import BaseModel, model_serializer, model_validator
@@ -112,7 +112,7 @@ class Data(BaseModel):
 
     def to_lc_message(
         self,
-    ) -> HumanMessage | SystemMessage:
+    ) -> BaseMessage:
         """
         Converts the Data to a BaseMessage.