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

refactor: Import BaseChatModel in OllamaModel.py

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-19 20:01:56 -03:00
commit 776ff92f57
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
src/backend/base/langflow/components/models/OllamaModel.py
View file

@ -1,5 +1,5 @@
from langchain_community.chat_models import ChatOllama 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.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
@ -176,7 +176,7 @@ class ChatOllamaComponent(LCModelComponent):
self.status = result self.status = result
return result return result
def build_model(self) -> BaseLanguageModel: def build_model(self) -> BaseLanguageModel | BaseChatModel:
# Mapping mirostat settings to their corresponding values # Mapping mirostat settings to their corresponding values
mirostat_options = {"Mirostat": 1, "Mirostat 2.0": 2} mirostat_options = {"Mirostat": 1, "Mirostat 2.0": 2}