Remove configurable prompt delimiters

Author: natoverse

graphrag/index/operations/extract_covariates/claim_extractor.py

@@ -20,12 +20,11 @@
 INPUT_ENTITY_SPEC_KEY = "entity_specs"
 INPUT_CLAIM_DESCRIPTION_KEY = "claim_description"
 INPUT_RESOLVED_ENTITIES_KEY = "resolved_entities"
-TUPLE_DELIMITER_KEY = "tuple_delimiter"
 RECORD_DELIMITER_KEY = "record_delimiter"
 COMPLETION_DELIMITER_KEY = "completion_delimiter"
-DEFAULT_TUPLE_DELIMITER = "<|>"
-DEFAULT_RECORD_DELIMITER = "##"
-DEFAULT_COMPLETION_DELIMITER = "<|COMPLETE|>"
+TUPLE_DELIMITER = "<|>"
+RECORD_DELIMITER = "##"
+COMPLETION_DELIMITER = "<|COMPLETE|>"
 logger = logging.getLogger(__name__)
 
 
@@ -71,20 +70,13 @@ async def __call__(
     ) -> ClaimExtractorResult:
         """Call method definition."""
         source_doc_map = {}
-
-        prompt_args = {
-            INPUT_ENTITY_SPEC_KEY: entity_spec,
-            INPUT_CLAIM_DESCRIPTION_KEY: claim_description,
-            TUPLE_DELIMITER_KEY: DEFAULT_TUPLE_DELIMITER,
-            RECORD_DELIMITER_KEY: DEFAULT_RECORD_DELIMITER,
-            COMPLETION_DELIMITER_KEY: DEFAULT_COMPLETION_DELIMITER,
-        }
-
         all_claims: list[dict] = []
         for doc_index, text in enumerate(texts):
             document_id = f"d{doc_index}"
             try:
-                claims = await self._process_document(prompt_args, text)
+                claims = await self._process_document(
+                    text, claim_description, entity_spec
+                )
                 all_claims += [
                     self._clean_claim(c, document_id, resolved_entities) for c in claims
                 ]
@@ -117,15 +109,18 @@ def _clean_claim(
         claim["subject_id"] = subject
         return claim
 
-    async def _process_document(self, prompt_args: dict, doc) -> list[dict]:
+    async def _process_document(
+        self, text: str, claim_description: str, entity_spec: dict
+    ) -> list[dict]:
         response = await self._model.achat(
             self._extraction_prompt.format(**{
-                INPUT_TEXT_KEY: doc,
-                **prompt_args,
+                INPUT_TEXT_KEY: text,
+                INPUT_CLAIM_DESCRIPTION_KEY: claim_description,
+                INPUT_ENTITY_SPEC_KEY: entity_spec,
             })
         )
         results = response.output.content or ""
-        claims = results.strip().removesuffix(DEFAULT_COMPLETION_DELIMITER)
+        claims = results.strip().removesuffix(COMPLETION_DELIMITER)
 
         # if gleanings are specified, enter a loop to extract more claims
         # there are two exit criteria: (a) we hit the configured max, (b) the model says there are no more claims
@@ -137,8 +132,8 @@ async def _process_document(self, prompt_args: dict, doc) -> list[dict]:
                     history=response.history,
                 )
                 extension = response.output.content or ""
-                claims += DEFAULT_RECORD_DELIMITER + extension.strip().removesuffix(
-                    DEFAULT_COMPLETION_DELIMITER
+                claims += RECORD_DELIMITER + extension.strip().removesuffix(
+                    COMPLETION_DELIMITER
                 )
 
                 # If this isn't the last loop, check to see if we should continue
@@ -164,18 +159,16 @@ def pull_field(index: int, fields: list[str]) -> str | None:
 
         result: list[dict[str, Any]] = []
         claims_values = (
-            claims.strip()
-            .removesuffix(DEFAULT_COMPLETION_DELIMITER)
-            .split(DEFAULT_RECORD_DELIMITER)
+            claims.strip().removesuffix(COMPLETION_DELIMITER).split(RECORD_DELIMITER)
         )
         for claim in claims_values:
             claim = claim.strip().removeprefix("(").removesuffix(")")
 
             # Ignore the completion delimiter
-            if claim == DEFAULT_COMPLETION_DELIMITER:
+            if claim == COMPLETION_DELIMITER:
                 continue
 
-            claim_fields = claim.split(DEFAULT_TUPLE_DELIMITER)
+            claim_fields = claim.split(TUPLE_DELIMITER)
             result.append({
                 "subject_id": pull_field(0, claim_fields),
                 "object_id": pull_field(1, claim_fields),

graphrag/index/operations/extract_graph/graph_extractor.py

@@ -22,12 +22,11 @@
 
 INPUT_TEXT_KEY = "input_text"
 RECORD_DELIMITER_KEY = "record_delimiter"
-TUPLE_DELIMITER_KEY = "tuple_delimiter"
 COMPLETION_DELIMITER_KEY = "completion_delimiter"
 ENTITY_TYPES_KEY = "entity_types"
-DEFAULT_TUPLE_DELIMITER = "<|>"
-DEFAULT_RECORD_DELIMITER = "##"
-DEFAULT_COMPLETION_DELIMITER = "<|COMPLETE|>"
+TUPLE_DELIMITER = "<|>"
+RECORD_DELIMITER = "##"
+COMPLETION_DELIMITER = "<|COMPLETE|>"
 DEFAULT_ENTITY_TYPES = ["organization", "person", "geo", "event"]
 
 logger = logging.getLogger(__name__)
@@ -72,18 +71,10 @@ async def __call__(
         all_records: dict[int, str] = {}
         source_doc_map: dict[int, str] = {}
 
-        # Wire defaults into the prompt variables
-        prompt_variables = {
-            ENTITY_TYPES_KEY: ",".join(entity_types),
-            TUPLE_DELIMITER_KEY: DEFAULT_TUPLE_DELIMITER,
-            RECORD_DELIMITER_KEY: DEFAULT_RECORD_DELIMITER,
-            COMPLETION_DELIMITER_KEY: DEFAULT_COMPLETION_DELIMITER,
-        }
-
         for doc_index, text in enumerate(texts):
             try:
                 # Invoke the entity extraction
-                result = await self._process_document(text, prompt_variables)
+                result = await self._process_document(text, entity_types)
                 source_doc_map[doc_index] = text
                 all_records[doc_index] = result
             except Exception as e:
@@ -99,22 +90,20 @@ async def __call__(
 
         output = await self._process_results(
             all_records,
-            DEFAULT_TUPLE_DELIMITER,
-            DEFAULT_RECORD_DELIMITER,
+            TUPLE_DELIMITER,
+            RECORD_DELIMITER,
         )
 
         return GraphExtractionResult(
             output=output,
             source_docs=source_doc_map,
         )
 
-    async def _process_document(
-        self, text: str, prompt_variables: dict[str, str]
-    ) -> str:
+    async def _process_document(self, text: str, entity_types: list[str]) -> str:
         response = await self._model.achat(
             self._extraction_prompt.format(**{
-                **prompt_variables,
                 INPUT_TEXT_KEY: text,
+                ENTITY_TYPES_KEY: ",".join(entity_types),
             }),
         )
         results = response.output.content or ""