From 72873c1b382bb4f8fd7ee42ced54fb28613df183 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Thu, 19 Sep 2024 12:34:25 +0200 Subject: [PATCH] refactoring of teh code --- examples/openai/smart_scraper_schema_openai.py | 2 +- scrapegraphai/nodes/generate_answer_node.py | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/examples/openai/smart_scraper_schema_openai.py b/examples/openai/smart_scraper_schema_openai.py index 0c1618d6..1df2be7b 100644 --- a/examples/openai/smart_scraper_schema_openai.py +++ b/examples/openai/smart_scraper_schema_openai.py @@ -29,7 +29,7 @@ openai_key = os.getenv("OPENAI_APIKEY") graph_config = { "llm": { - "api_key":openai_key, + "api_key": openai_key, "model": "openai/gpt-4o-mini", }, "verbose": True, diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index b0c102e1..3d1c0481 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -82,18 +82,16 @@ class GenerateAnswerNode(BaseNode): self.logger.info(f"--- Executing {self.node_name} Node ---") - input_keys = self.get_input_keys(state) - + input_keys = self.get_input_keys(state) input_data = [state[key] for key in input_keys] user_prompt = input_data[0] doc = input_data[1] if self.node_config.get("schema", None) is not None: - + if isinstance(self.llm_model, (ChatOpenAI, ChatMistralAI)): self.llm_model = self.llm_model.with_structured_output( - schema = self.node_config["schema"]) # json schema works only on specific models - + schema = self.node_config["schema"]) output_parser = get_structured_output_parser(self.node_config["schema"]) format_instructions = "NA" else: @@ -120,7 +118,7 @@ class GenerateAnswerNode(BaseNode): if self.additional_info is not None: template_no_chunks_prompt = self.additional_info + template_no_chunks_prompt template_chunks_prompt = self.additional_info + template_chunks_prompt - template_merge_prompt = self.additional_info + template_merge_prompt + template_merge_prompt = self.additional_info + template_merge_prompt if len(doc) == 1: prompt = PromptTemplate(