From 43058045ebc082ea4f2007a5ca6d84028641c251 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 20 Jun 2024 18:09:10 -0300 Subject: [PATCH] refactor: Update LCModelComponent to use BaseMessage instead of HumanMessage and SystemMessage --- src/backend/base/langflow/base/models/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/base/models/model.py b/src/backend/base/langflow/base/models/model.py index 810fa092c..1e4152b28 100644 --- a/src/backend/base/langflow/base/models/model.py +++ b/src/backend/base/langflow/base/models/model.py @@ -3,7 +3,7 @@ import warnings from typing import Optional, Union from langchain_core.language_models.llms import LLM -from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage from langflow.custom import Component from langflow.field_typing import LanguageModel @@ -120,7 +120,7 @@ class LCModelComponent(Component): def get_chat_result( self, runnable: LanguageModel, stream: bool, input_value: str | Message, system_message: Optional[str] = None ): - messages: list[Union[HumanMessage, SystemMessage]] = [] + messages: list[Union[BaseMessage]] = [] if not input_value and not system_message: raise ValueError("The message you want to send to the model is empty.") if system_message: