feat(operate): 添加实体类型配置并优化提示生成
Browse files- 在全局配置中添加 entity_types 参数,用于自定义实体类型
- 在生成实体提取和关系提取的提示时,使用配置的实体类型替代默认值
- 优化了提示生成逻辑,提高了代码的可配置性和灵活性
- 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(
|
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(
|
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 |
)
|