From ab00f23d859c64995ccfe329b24379cf3c14d73c Mon Sep 17 00:00:00 2001 From: Marco Perini Date: Wed, 12 Jun 2024 01:40:49 +0200 Subject: [PATCH] fix(node): fixed generate answer node pydantic schema --- scrapegraphai/nodes/generate_answer_node.py | 23 ++++----------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index b5ec4a3d..c6b8c388 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -93,35 +93,20 @@ 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.get("schema", None) is None and len(doc) == 1: + if 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.get("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.get("schema", None) - }) - elif self.node_config.get("schema", None) is None and len(doc) > 1: + + else: 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.get("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.get("schema", None)}) # Dynamically name the chains based on their index chain_name = f"chunk{i+1}" @@ -147,4 +132,4 @@ class GenerateAnswerNode(BaseNode): # Update the state with the generated answer state.update({self.output[0]: answer}) - return state + return state \ No newline at end of file