feat - in local and global search system prompts, add the support to …

…enable/disable repeated instructions
ksachdeva · Oct 8, 2024 · 69d9aa9 · 69d9aa9
1 parent 1868509
commit 69d9aa9
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 4 deletions.
diff --git a/examples/simple-app/app/query.py b/examples/simple-app/app/query.py
@@ -72,6 +72,10 @@ def global_search(
         None, help="Context window size for ollama model"
     ),
     show_references: bool = typer.Option(True, help="Show references in the output"),  # noqa: FBT001, FBT003
+    repeat_instructions: bool = typer.Option(  # noqa: FBT001
+        True,  # noqa: FBT003
+        help="Repeat instructions in the prompt",
+    ),
     enable_langsmith: bool = typer.Option(False, help="Enable Langsmith"),  # noqa: FBT001, FBT003
 ):
     if enable_langsmith:
@@ -94,6 +98,7 @@ def global_search(
                 "Not Provided" if ollama_num_context is None else ollama_num_context,
             ],
             ["Show References", str(show_references)],
+            ["Repeat Instructions In Prompt", str(repeat_instructions)],
         ]
     )
 
@@ -113,14 +118,17 @@ def global_search(
             cache_dir,
             ollama_num_context=ollama_num_context,
         ),
-        prompt_builder=KeyPointsGeneratorPromptBuilder(show_references=show_references),
+        prompt_builder=KeyPointsGeneratorPromptBuilder(
+            show_references=show_references, repeat_instructions=repeat_instructions
+        ),
         context_builder=report_context_builder,
     )
 
     kp_aggregator = KeyPointsAggregator(
         llm=make_llm_instance(llm_type, llm_model, cache_dir),
         prompt_builder=KeyPointsAggregatorPromptBuilder(
-            show_references=show_references
+            show_references=show_references,
+            repeat_instructions=repeat_instructions,
         ),
         context_builder=KeyPointsContextBuilder(
             token_counter=TiktokenCounter(),
@@ -157,6 +165,10 @@ def local_search(
         None, help="Context window size for ollama model"
     ),
     show_references: bool = typer.Option(True, help="Show references in the output"),  # noqa: FBT001, FBT003
+    repeat_instructions: bool = typer.Option(  # noqa: FBT001
+        True,  # noqa: FBT003
+        help="Repeat instructions in the prompt",
+    ),
     enable_langsmith: bool = typer.Option(False, help="Enable Langsmith"),  # noqa: FBT001, FBT003
 ):
     if enable_langsmith:
@@ -183,6 +195,7 @@ def local_search(
                 "Not Provided" if ollama_num_context is None else ollama_num_context,
             ],
             ["Show References", str(show_references)],
+            ["Repeat Instructions In Prompt", str(repeat_instructions)],
         ]
     )
 
@@ -231,7 +244,10 @@ def local_search(
     )
 
     local_search = LocalSearch(
-        prompt_builder=LocalSearchPromptBuilder(show_references=show_references),
+        prompt_builder=LocalSearchPromptBuilder(
+            show_references=show_references,
+            repeat_instructions=repeat_instructions,
+        ),
         llm=make_llm_instance(
             llm_type,
             llm_model,

diff --git a/src/langchain_graphrag/query/global_search/key_points_aggregator/_system_prompt.py b/src/langchain_graphrag/query/global_search/key_points_aggregator/_system_prompt.py
@@ -39,10 +39,14 @@
 
 {{response_type}}
 
+Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+
 ---Analyst Reports---
 
 {{report_data}}
 
+{{#repeat_instructions}}
+
 ---Goal---
 
 Generate a response of the target length and format that responds to the user's question, summarize all the reports from multiple analysts who focused on different parts of the dataset.
@@ -76,4 +80,5 @@
 {{response_type}}
 
 Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+{{/repeat_instructions}}
 """
diff --git a/src/langchain_graphrag/query/global_search/key_points_aggregator/prompt_builder.py b/src/langchain_graphrag/query/global_search/key_points_aggregator/prompt_builder.py
@@ -20,6 +20,7 @@ def __init__(
         system_prompt: str | None = None,
         system_prompt_path: Path | None = None,
         show_references: bool = True,
+        repeat_instructions: bool = True,
     ):
         self._system_prompt: str | None
         if system_prompt is None and system_prompt_path is None:
@@ -29,6 +30,7 @@ def __init__(
 
         self._system_prompt_path = system_prompt_path
         self._show_references = show_references
+        self._repeat_instructions = repeat_instructions
 
     def build(self) -> tuple[BasePromptTemplate, BaseOutputParser]:
         if self._system_prompt_path:
@@ -42,6 +44,7 @@ def build(self) -> tuple[BasePromptTemplate, BaseOutputParser]:
             partial_variables=dict(
                 response_type="Multiple Paragraphs",
                 show_references=self._show_references,
+                repeat_instructions=self._repeat_instructions,
             ),
             template_format="mustache",
         )

diff --git a/src/langchain_graphrag/query/global_search/key_points_generator/_system_prompt.py b/src/langchain_graphrag/query/global_search/key_points_generator/_system_prompt.py
@@ -47,6 +47,8 @@
 
 {{context_data}}
 
+{{#repeat_instructions}}
+
 ---Goal---
 
 Generate a response consisting of a list of key points that responds to the user's question, summarizing all relevant information in the input data tables.
@@ -83,4 +85,5 @@
         {"description": "Description of point 2 {{#show_references}}[Data: Reports (report ids)]{{/show_references}}", "score": score_value}
     ]
 }
+{{/repeat_instructions}}
 """
diff --git a/src/langchain_graphrag/query/global_search/key_points_generator/prompt_builder.py b/src/langchain_graphrag/query/global_search/key_points_generator/prompt_builder.py
@@ -20,6 +20,7 @@ def __init__(
         system_prompt: str | None = None,
         system_prompt_path: Path | None = None,
         show_references: bool = True,
+        repeat_instructions: bool = True,
     ):
         self._system_prompt: str | None
         if system_prompt is None and system_prompt_path is None:
@@ -29,6 +30,7 @@ def __init__(
 
         self._system_prompt_path = system_prompt_path
         self._show_references = show_references
+        self._repeat_instructions = repeat_instructions
 
     def build(self) -> tuple[BasePromptTemplate, BaseOutputParser]:
         if self._system_prompt_path:
@@ -40,7 +42,10 @@ def build(self) -> tuple[BasePromptTemplate, BaseOutputParser]:
         system_template = SystemMessagePromptTemplate.from_template(
             prompt,
             template_format="mustache",
-            partial_variables=dict(show_references=self._show_references),
+            partial_variables=dict(
+                show_references=self._show_references,
+                repeat_instructions=self._repeat_instructions,
+            ),
         )
 
         template = ChatPromptTemplate(

diff --git a/src/langchain_graphrag/query/local_search/_system_prompt.py b/src/langchain_graphrag/query/local_search/_system_prompt.py
@@ -33,10 +33,14 @@
 
 {{response_type}}
 
+Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+
 ---Data tables---
 
 {{context_data}}
 
+{{#repeat_instructions}}
+
 ---Goal---
 
 Generate a response of the target length and format that responds to the user's question, summarizing all information in the input data tables appropriate for the response length and format, and incorporating any relevant general knowledge.
@@ -64,4 +68,6 @@
 {{response_type}}
 
 Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+
+{{/repeat_instructions}}
 """
diff --git a/src/langchain_graphrag/query/local_search/prompt_builder.py b/src/langchain_graphrag/query/local_search/prompt_builder.py
@@ -20,6 +20,7 @@ def __init__(
         system_prompt: str | None = None,
         system_prompt_path: Path | None = None,
         show_references: bool = True,
+        repeat_instructions: bool = True,
     ):
         self._system_prompt: str | None
         if system_prompt is None and system_prompt_path is None:
@@ -29,6 +30,7 @@ def __init__(
 
         self._system_prompt_path = system_prompt_path
         self._show_references = show_references
+        self._repeat_instructions = repeat_instructions
 
     def build(self) -> tuple[BasePromptTemplate, BaseOutputParser]:
         if self._system_prompt_path:
@@ -42,6 +44,7 @@ def build(self) -> tuple[BasePromptTemplate, BaseOutputParser]:
             partial_variables=dict(
                 response_type="Multiple Paragraphs",
                 show_references=self._show_references,
+                repeat_instructions=self._repeat_instructions,
             ),
             template_format="mustache",
         )