From 51fbd70508e15c92ca41e30cb4a9f6161ffccec5 Mon Sep 17 00:00:00 2001 From: AmeyaBadve Date: Wed, 14 Oct 2026 00:57:22 -0400 Subject: [PATCH 1/2] fix: include user question in chunk extraction prompts --- .../prompts/generate_answer_node_prompts.py | 2 ++ tests/test_chunk_prompt_question.py | 30 +++++++++++++++++++ 2 files changed, 32 insertions(+) create mode 100644 tests/test_chunk_prompt_question.py diff --git a/scrapegraphai/prompts/generate_answer_node_prompts.py b/scrapegraphai/prompts/generate_answer_node_prompts.py index 91e49c250..9af5babee 100644 --- a/scrapegraphai/prompts/generate_answer_node_prompts.py +++ b/scrapegraphai/prompts/generate_answer_node_prompts.py @@ -13,6 +13,7 @@ and things that will invalidate the dictionary. \n Do not start the response with ```json because it will invalidate the postprocessing. \n OUTPUT INSTRUCTIONS: {format_instructions}\n +USER QUESTION: {question} Content of {chunk_id}: {content}. \n """ @@ -56,6 +57,7 @@ and things that will invalidate the dictionary. \n Do not start the response with ```json because it will invalidate the postprocessing. \n OUTPUT INSTRUCTIONS: {format_instructions}\n +USER QUESTION: {question} Content of {chunk_id}: {content}. \n """ diff --git a/tests/test_chunk_prompt_question.py b/tests/test_chunk_prompt_question.py new file mode 100644 index 000000000..6dcfeb236 --- /dev/null +++ b/tests/test_chunk_prompt_question.py @@ -0,0 +1,30 @@ +import pytest +from langchain_core.prompts import PromptTemplate +from scrapegraphai import prompts + + +@pytest.mark.parametrize( + "template_name", + ["TEMPLATE_CHUNKS_MD", "TEMPLATE_CHUNKS"], +) +def test_chunk_prompts_include_user_question(template_name): + """Ensure each chunk receives the user's extraction question.""" + from langchain_core.prompts import PromptTemplate + from scrapegraphai import prompts + + template = getattr(prompts, template_name) + + prompt = PromptTemplate( + template=template, + input_variables=["question"], + partial_variables={ + "content": "Django and Flask are Python web frameworks.", + "chunk_id": 1, + "format_instructions": "Return valid JSON.", + }, + ) + + rendered = prompt.format(question="Which Python frameworks are mentioned?") + + assert "USER QUESTION: Which Python frameworks are mentioned?" in rendered + assert "Django and Flask" in rendered From f424c56a9021c880223ac06fcc40cae4a603b849 Mon Sep 17 00:00:00 2001 From: AmeyaBadve Date: Wed, 14 Oct 2026 01:12:05 -0400 Subject: [PATCH 2/2] fix: pass user question to K-level chunk prompts --- .../nodes/generate_answer_node_k_level.py | 5 ++- tests/test_chunk_prompt_question.py | 37 +++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/scrapegraphai/nodes/generate_answer_node_k_level.py b/scrapegraphai/nodes/generate_answer_node_k_level.py index 0a51866a0..173018beb 100644 --- a/scrapegraphai/nodes/generate_answer_node_k_level.py +++ b/scrapegraphai/nodes/generate_answer_node_k_level.py @@ -149,17 +149,18 @@ def execute(self, state: dict) -> dict: ): prompt = PromptTemplate( template=template_chunks_prompt, - input_variables=["format_instructions"], + input_variables=["question"], partial_variables={ "content": chunk.get("document"), "chunk_id": i + 1, + "format_instructions": format_instructions, }, ) chain_name = f"chunk{i + 1}" chains_dict[chain_name] = prompt | self.llm_model async_runner = RunnableParallel(**chains_dict) - batch_results = async_runner.invoke({"format_instructions": user_prompt}) + batch_results = async_runner.invoke({"question": user_prompt}) merge_prompt = PromptTemplate( template=template_merge_prompt, diff --git a/tests/test_chunk_prompt_question.py b/tests/test_chunk_prompt_question.py index 6dcfeb236..9764b2031 100644 --- a/tests/test_chunk_prompt_question.py +++ b/tests/test_chunk_prompt_question.py @@ -28,3 +28,40 @@ def test_chunk_prompts_include_user_question(template_name): assert "USER QUESTION: Which Python frameworks are mentioned?" in rendered assert "Django and Flask" in rendered + + +def test_k_level_chunk_receives_question(monkeypatch): + """Verify K-level chunk processing includes the question and JSON instructions.""" + from langchain_core.runnables import RunnableLambda + from scrapegraphai.nodes.generate_answer_node_k_level import GenerateAnswerNodeKLevel + + captured_prompts = [] + + def fake_llm(prompt): + captured_prompts.append(prompt.to_string()) + return '{"content": "Django"}' + + class FakeResult: + id = 1 + score = 0.9 + + class FakeDB: + def query(self, collection_name, query_text): + return [FakeResult()] + + node = GenerateAnswerNodeKLevel( + "user_prompt & docs & vectorial_db", + ["answer"], + node_config={"llm_model": RunnableLambda(fake_llm), "verbose": False}, + ) + + node.execute({ + "user_prompt": "Which Python frameworks are mentioned?", + "docs": [{"document": "Django is a Python framework."}], + "vectorial_db": FakeDB(), + }) + + assert captured_prompts + assert "USER QUESTION: Which Python frameworks are mentioned?" in captured_prompts[0] + assert "OUTPUT INSTRUCTIONS:" in captured_prompts[0] + assert "OUTPUT INSTRUCTIONS: Return a JSON object." in captured_prompts[0]