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/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..9764b2031 --- /dev/null +++ b/tests/test_chunk_prompt_question.py @@ -0,0 +1,67 @@ +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 + + +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]