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)