jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: Update CohereModel to import BaseChatModel from langchain_core

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-19 20:02:10 -03:00
commit 7809fba2aa
1 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/backend/base/langflow/components/models/CohereModel.py
View file

@ -1,10 +1,11 @@
from langchain_cohere import ChatCohere from langchain_cohere import ChatCohere
from langchain_core.language_models.chat_models import BaseChatModel
from pydantic.v1 import SecretStr from pydantic.v1 import SecretStr
from langflow.base.constants import STREAM_INFO_TEXT from langflow.base.constants import STREAM_INFO_TEXT
from langflow.base.models.model import LCModelComponent from langflow.base.models.model import LCModelComponent
from langflow.field_typing import BaseLanguageModel, Text 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): class CohereComponent(LCModelComponent):
@ -45,7 +46,7 @@ class CohereComponent(LCModelComponent):
self.status = result self.status = result
return result return result
def build_model(self) -> BaseLanguageModel: def build_model(self) -> BaseLanguageModel | BaseChatModel:
cohere_api_key = self.cohere_api_key cohere_api_key = self.cohere_api_key
temperature = self.temperature temperature = self.temperature