Magicyuan commited on
Commit
b6f2187
·
1 Parent(s): 37c0b15

feat(operate): 添加实体类型配置并优化提示生成

Browse files

- 在全局配置中添加 entity_types 参数,用于自定义实体类型
- 在生成实体提取和关系提取的提示时,使用配置的实体类型替代默认值
- 优化了提示生成逻辑,提高了代码的可配置性和灵活性

Files changed (1) hide show
  1. lightrag/operate.py +5 -2
lightrag/operate.py CHANGED
@@ -260,6 +260,9 @@ async def extract_entities(
260
  language = global_config["addon_params"].get(
261
  "language", PROMPTS["DEFAULT_LANGUAGE"]
262
  )
 
 
 
263
  example_number = global_config["addon_params"].get("example_number", None)
264
  if example_number and example_number < len(PROMPTS["entity_extraction_examples"]):
265
  examples = "\n".join(
@@ -272,7 +275,7 @@ async def extract_entities(
272
  tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
273
  record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
274
  completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
275
- entity_types=",".join(PROMPTS["DEFAULT_ENTITY_TYPES"]),
276
  language=language,
277
  )
278
  # add example's format
@@ -283,7 +286,7 @@ async def extract_entities(
283
  tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
284
  record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
285
  completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
286
- entity_types=",".join(PROMPTS["DEFAULT_ENTITY_TYPES"]),
287
  examples=examples,
288
  language=language,
289
  )
 
260
  language = global_config["addon_params"].get(
261
  "language", PROMPTS["DEFAULT_LANGUAGE"]
262
  )
263
+ entity_types = global_config["addon_params"].get(
264
+ "entity_types", PROMPTS["DEFAULT_ENTITY_TYPES"]
265
+ )
266
  example_number = global_config["addon_params"].get("example_number", None)
267
  if example_number and example_number < len(PROMPTS["entity_extraction_examples"]):
268
  examples = "\n".join(
 
275
  tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
276
  record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
277
  completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
278
+ entity_types=",".join(entity_types),
279
  language=language,
280
  )
281
  # add example's format
 
286
  tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
287
  record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
288
  completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
289
+ entity_types=",".join(entity_types),
290
  examples=examples,
291
  language=language,
292
  )