From c8d556da4e4b8730c6c35f1d448270b8e26923f2 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Mon, 3 Jun 2024 21:49:34 +0200 Subject: [PATCH] feat: fix an if --- scrapegraphai/nodes/generate_answer_node.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index c57de035..4b8c5c36 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -12,7 +12,7 @@ from langchain_core.runnables import RunnableParallel from tqdm import tqdm from ..utils.logging import get_logger -from ..models import Ollama, Groq +from ..models import Ollama, Groq, OpenAI # Imports from the library from .base_node import BaseNode from ..helpers import template_chunks, template_no_chunks, template_merge, template_chunks_with_schema, template_no_chunks_with_schema @@ -46,7 +46,7 @@ class GenerateAnswerNode(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"], Ollama) or isinstance(node_config["llm_model"], OpenAI): self.llm_model.format="json" self.verbose = ( True if node_config is None else node_config.get("verbose", False)