From 1d38ed146afae95dae1f35ac51180a1882bf8a29 Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Wed, 5 Jun 2024 09:17:29 +0200
Subject: [PATCH] fix: bug on generate_answer_node

---
 scrapegraphai/nodes/generate_answer_node.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py
index 26a2ed66..22461508 100644
--- a/scrapegraphai/nodes/generate_answer_node.py
+++ b/scrapegraphai/nodes/generate_answer_node.py
@@ -83,35 +83,35 @@ class GenerateAnswerNode(BaseNode):
 
         # Use tqdm to add progress bar
         for i, chunk in enumerate(tqdm(doc, desc="Processing chunks", disable=not self.verbose)):
-            if self.node_config["schema"] is None and len(doc) == 1:
+            if self.node_config("schema", None) is None and len(doc) == 1:
                 prompt = PromptTemplate(
                     template=template_no_chunks,
                     input_variables=["question"],
                     partial_variables={"context": chunk.page_content,
                                        "format_instructions": format_instructions})
-            elif self.node_config["schema"] is not None and len(doc) == 1:
+            elif self.node_config("schema", None) is not None and len(doc) == 1:
                  prompt = PromptTemplate(
                     template=template_no_chunks_with_schema,
                     input_variables=["question"],
                     partial_variables={"context": chunk.page_content,
                                        "format_instructions": format_instructions,
-                                       "schema": self.node_config["schema"]
+                                       "schema": self.node_config("schema", None)
                                        })
-            elif self.node_config["schema"] is None and len(doc) > 1:
+            elif self.node_config("schema", None) is None and len(doc) > 1:
                 prompt = PromptTemplate(
                     template=template_chunks,
                     input_variables=["question"],
                     partial_variables={"context": chunk.page_content,
                                         "chunk_id": i + 1,
                                         "format_instructions": format_instructions})
-            elif self.node_config["schema"] is not None and len(doc) > 1:
+            elif self.node_config("schema", None) is not None and len(doc) > 1:
                 prompt = PromptTemplate(
                     template=template_chunks_with_schema,
                     input_variables=["question"],
                     partial_variables={"context": chunk.page_content,
                                         "chunk_id": i + 1,
                                         "format_instructions": format_instructions,
-                                        "schema": self.node_config["schema"]})
+                                        "schema": self.node_config("schema", None)})
 
             # Dynamically name the chains based on their index
             chain_name = f"chunk{i+1}"