From 776ff92f57794ec0e6661231815c84741fc36515 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 19 Jun 2024 20:01:56 -0300 Subject: [PATCH] refactor: Import BaseChatModel in OllamaModel.py --- src/backend/base/langflow/components/models/OllamaModel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/components/models/OllamaModel.py b/src/backend/base/langflow/components/models/OllamaModel.py index 3b2872951..4173b9e13 100644 --- a/src/backend/base/langflow/components/models/OllamaModel.py +++ b/src/backend/base/langflow/components/models/OllamaModel.py @@ -1,5 +1,5 @@ from langchain_community.chat_models import ChatOllama - +from langchain_core.language_models.chat_models import BaseChatModel from langflow.base.constants import STREAM_INFO_TEXT from langflow.base.models.model import LCModelComponent from langflow.field_typing import BaseLanguageModel, Text @@ -176,7 +176,7 @@ class ChatOllamaComponent(LCModelComponent): self.status = result return result - def build_model(self) -> BaseLanguageModel: + def build_model(self) -> BaseLanguageModel | BaseChatModel: # Mapping mirostat settings to their corresponding values mirostat_options = {"Mirostat": 1, "Mirostat 2.0": 2}