From 7809fba2aa952d0f01709806d7ccd138da1f30f9 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 19 Jun 2024 20:02:10 -0300 Subject: [PATCH] refactor: Update CohereModel to import BaseChatModel from langchain_core --- src/backend/base/langflow/components/models/CohereModel.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/components/models/CohereModel.py b/src/backend/base/langflow/components/models/CohereModel.py index e22a0c98e..6f1baa424 100644 --- a/src/backend/base/langflow/components/models/CohereModel.py +++ b/src/backend/base/langflow/components/models/CohereModel.py @@ -1,10 +1,11 @@ from langchain_cohere import ChatCohere +from langchain_core.language_models.chat_models import BaseChatModel from pydantic.v1 import SecretStr from langflow.base.constants import STREAM_INFO_TEXT from langflow.base.models.model import LCModelComponent from langflow.field_typing import BaseLanguageModel, Text -from langflow.io import BoolInput, FloatInput, IntInput, MessageInput, Output, SecretStrInput, StrInput +from langflow.io import BoolInput, FloatInput, MessageInput, Output, SecretStrInput, StrInput class CohereComponent(LCModelComponent): @@ -45,7 +46,7 @@ class CohereComponent(LCModelComponent): self.status = result return result - def build_model(self) -> BaseLanguageModel: + def build_model(self) -> BaseLanguageModel | BaseChatModel: cohere_api_key = self.cohere_api_key temperature = self.temperature