Magicyuan commited on
Commit
4017b66
·
1 Parent(s): 0165c41

feat(llm, prompt):增加日志输出并扩展实体类型

Browse files

- 在 llm.py 中添加了日志输出,用于调试和记录 LLM 查询输入
- 在 prompt.py 中增加了 "category" 实体类型,扩展了实体提取的范围

Files changed (2) hide show
  1. lightrag/llm.py +6 -1
  2. lightrag/prompt.py +1 -1
lightrag/llm.py CHANGED
@@ -29,7 +29,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
29
  from .utils import (
30
  wrap_embedding_func_with_attrs,
31
  locate_json_string_body_from_string,
32
- safe_unicode_decode,
33
  )
34
 
35
  import sys
@@ -69,6 +69,11 @@ async def openai_complete_if_cache(
69
  messages.extend(history_messages)
70
  messages.append({"role": "user", "content": prompt})
71
 
 
 
 
 
 
72
  if "response_format" in kwargs:
73
  response = await openai_async_client.beta.chat.completions.parse(
74
  model=model, messages=messages, **kwargs
 
29
  from .utils import (
30
  wrap_embedding_func_with_attrs,
31
  locate_json_string_body_from_string,
32
+ safe_unicode_decode, logger,
33
  )
34
 
35
  import sys
 
69
  messages.extend(history_messages)
70
  messages.append({"role": "user", "content": prompt})
71
 
72
+ # 添加日志输出
73
+ logger.debug("===== Query Input to LLM =====")
74
+ logger.debug(f"Query: {prompt}")
75
+ logger.debug(f"System prompt: {system_prompt}")
76
+ logger.debug("Full context:")
77
  if "response_format" in kwargs:
78
  response = await openai_async_client.beta.chat.completions.parse(
79
  model=model, messages=messages, **kwargs
lightrag/prompt.py CHANGED
@@ -8,7 +8,7 @@ PROMPTS["DEFAULT_RECORD_DELIMITER"] = "##"
8
  PROMPTS["DEFAULT_COMPLETION_DELIMITER"] = "<|COMPLETE|>"
9
  PROMPTS["process_tickers"] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
10
 
11
- PROMPTS["DEFAULT_ENTITY_TYPES"] = ["organization", "person", "geo", "event"]
12
 
13
  PROMPTS["entity_extraction"] = """-Goal-
14
  Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text and all relationships among the identified entities.
 
8
  PROMPTS["DEFAULT_COMPLETION_DELIMITER"] = "<|COMPLETE|>"
9
  PROMPTS["process_tickers"] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
10
 
11
+ PROMPTS["DEFAULT_ENTITY_TYPES"] = ["organization", "person", "geo", "event", "category"]
12
 
13
  PROMPTS["entity_extraction"] = """-Goal-
14
  Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text and all relationships among the identified entities.