From 5fe694b6b4545a5091d16110318b992acfca4f58 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Sun, 30 Jun 2024 18:10:00 +0200 Subject: [PATCH] feat: improve md prompt recognition --- scrapegraphai/graphs/markdown_scraper_graph.py | 1 + scrapegraphai/nodes/generate_answer_node.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/scrapegraphai/graphs/markdown_scraper_graph.py b/scrapegraphai/graphs/markdown_scraper_graph.py index 655aee94..7fb3f10f 100644 --- a/scrapegraphai/graphs/markdown_scraper_graph.py +++ b/scrapegraphai/graphs/markdown_scraper_graph.py @@ -77,6 +77,7 @@ class MDScraperGraph(AbstractGraph): node_config={ "llm_model": self.llm_model, "schema": self.schema, + "is_md_scraper": True } ) diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index 476421f0..b2ea63ee 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -53,6 +53,9 @@ class GenerateAnswerNode(BaseNode): self.script_creator = ( False if node_config is None else node_config.get("script_creator", False) ) + self.is_md_scraper = ( + False if node_config is None else node_config.get("is_md_scraper", False) + ) def execute(self, state: dict) -> dict: @@ -89,7 +92,7 @@ class GenerateAnswerNode(BaseNode): format_instructions = output_parser.get_format_instructions() - if isinstance(self.llm_model, OpenAI) and not self.script_creator or self.force and not self.script_creator: + if isinstance(self.llm_model, OpenAI) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper: template_no_chunks_prompt = template_no_chunks_md template_chunks_prompt = template_chunks_md template_merge_prompt = template_merge_md