From d32a35d8fdc154a5f8e7e4773608a60324f12294 Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Fri, 28 Mar 2025 15:00:24 +0800 Subject: [PATCH] Fix entity_types. Close #6287 and #6608 (#6632) ### What problem does this PR solve? Fix entity_types. Close #6287 and #6608 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) --- graphrag/general/graph_extractor.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/graphrag/general/graph_extractor.py b/graphrag/general/graph_extractor.py index b2a3948a4..75e023698 100644 --- a/graphrag/general/graph_extractor.py +++ b/graphrag/general/graph_extractor.py @@ -11,7 +11,7 @@ from dataclasses import dataclass import tiktoken import trio -from graphrag.general.extractor import Extractor, ENTITY_EXTRACTION_MAX_GLEANINGS, DEFAULT_ENTITY_TYPES +from graphrag.general.extractor import Extractor, ENTITY_EXTRACTION_MAX_GLEANINGS from graphrag.general.graph_prompt import GRAPH_EXTRACTION_PROMPT, CONTINUE_PROMPT, LOOP_PROMPT from graphrag.utils import ErrorHandlerFn, perform_variable_replacements, chat_limiter, split_string_by_multi_markers from rag.llm.chat_model import Base as CompletionLLM @@ -91,11 +91,10 @@ class GraphExtractor(Extractor): # Wire defaults into the prompt variables self._prompt_variables = { - "entity_types": entity_types, self._tuple_delimiter_key: DEFAULT_TUPLE_DELIMITER, self._record_delimiter_key: DEFAULT_RECORD_DELIMITER, self._completion_delimiter_key: DEFAULT_COMPLETION_DELIMITER, - self._entity_types_key: ",".join(DEFAULT_ENTITY_TYPES), + self._entity_types_key: entity_types, } async def _process_single_content(self, chunk_key_dp: tuple[str, str], chunk_seq: int, num_chunks: int, out_results):