refactor: remove LangChain wrappers for Ollama

This commit is contained in:
Federico Aguzzi 2024-07-29 17:49:21 +02:00
parent 2c5f934f10
commit 25066b2bc5
2 changed files with 4 additions and 4 deletions

View File

@ -10,7 +10,7 @@ from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from tqdm import tqdm
from ..models import Ollama
from langchain_community.chat_models import ChatOllama
# Imports from the library
from .base_node import BaseNode
from ..helpers.generate_answer_node_omni_prompts import template_no_chunk_omni, template_chunks_omni, template_merge_omni
@ -44,7 +44,7 @@ class GenerateAnswerOmniNode(BaseNode):
super().__init__(node_name, "node", input, output, 3, node_config)
self.llm_model = node_config["llm_model"]
if isinstance(node_config["llm_model"], Ollama):
if isinstance(node_config["llm_model"], ChatOllama):
self.llm_model.format="json"
self.verbose = (

View File

@ -10,7 +10,7 @@ from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from tqdm import tqdm
from ..models import Ollama
from langchain_community.chat_models import ChatOllama
from ..utils.logging import get_logger
# Imports from the library
@ -59,7 +59,7 @@ class GenerateAnswerPDFNode(BaseNode):
super().__init__(node_name, "node", input, output, 2, node_config)
self.llm_model = node_config["llm_model"]
if isinstance(node_config["llm_model"], Ollama):
if isinstance(node_config["llm_model"], ChatOllama):
self.llm_model.format="json"
self.verbose = (