update prompt var usage

microsoft · Oct 25, 2024 · 4795495 · 4795495
1 parent 9d607a1
commit 4795495
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 12 deletions.
diff --git a/graphrag/index/graph/extractors/claims/claim_extractor.py b/graphrag/index/graph/extractors/claims/claim_extractor.py
@@ -167,11 +167,10 @@ async def _process_document(
         )
 
         response = await self._llm(
-            self._extraction_prompt,
-            variables={
+            self._extraction_prompt.format(**{
                 self._input_text_key: doc,
                 **prompt_args,
-            },
+            })
         )
         results = response.output.content or ""
         claims = results.strip().removesuffix(completion_delimiter)

diff --git a/graphrag/index/graph/extractors/community_reports/community_reports_extractor.py b/graphrag/index/graph/extractors/community_reports/community_reports_extractor.py
@@ -75,10 +75,11 @@ async def __call__(self, inputs: dict[str, Any]):
         output = None
         try:
             response = await self._llm(
-                self._extraction_prompt,
+                self._extraction_prompt.format(**{
+                    self._input_text_key: inputs[self._input_text_key]
+                }),
                 json=True,
                 name="create_community_report",
-                variables={self._input_text_key: inputs[self._input_text_key]},
                 json_model=CommunityReportResponse,
                 model_parameters={"max_tokens": self._max_report_length},
             )

diff --git a/graphrag/index/graph/extractors/graph/graph_extractor.py b/graphrag/index/graph/extractors/graph/graph_extractor.py
@@ -149,11 +149,10 @@ async def _process_document(
         self, text: str, prompt_variables: dict[str, str]
     ) -> str:
         response = await self._llm(
-            self._extraction_prompt,
-            variables={
+            self._extraction_prompt.format(**{
                 **prompt_variables,
                 self._input_text_key: text,
-            },
+            }),
         )
         results = response.output.content or ""
 

diff --git a/graphrag/index/graph/extractors/summarize/description_summary_extractor.py b/graphrag/index/graph/extractors/summarize/description_summary_extractor.py
@@ -128,14 +128,13 @@ async def _summarize_descriptions_with_llm(
     ):
         """Summarize descriptions using the LLM."""
         response = await self._llm(
-            self._summarization_prompt,
-            name="summarize",
-            variables={
+            self._summarization_prompt.format(**{
                 self._entity_name_key: json.dumps(items),
                 self._input_descriptions_key: json.dumps(
                     sorted(descriptions), ensure_ascii=False
                 ),
-            },
+            }),
+            name="summarize",
             model_parameters={"max_tokens": self._max_summary_length},
         )
         # Calculate result