gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 在 llm.py 中添加了日志输出，用于调试和记录 LLM 查询输入
- 在 prompt.py 中增加了 "category" 实体类型，扩展了实体提取的范围
This commit is contained in:
Magic_yuan
2024-12-11 12:45:10 +08:00
parent 316c4df949
commit 0a41cc8a9a
2 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
lightrag/llm.py
View File

@@ -29,7 +29,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
from .utils import ( from .utils import (
wrap_embedding_func_with_attrs, wrap_embedding_func_with_attrs,
locate_json_string_body_from_string, locate_json_string_body_from_string,
safe_unicode_decode, safe_unicode_decode, logger,
) )
import sys import sys
@@ -69,6 +69,11 @@ async def openai_complete_if_cache(
messages.extend(history_messages) messages.extend(history_messages)
messages.append({"role": "user", "content": prompt}) messages.append({"role": "user", "content": prompt})
# 添加日志输出
logger.debug("===== Query Input to LLM =====")
logger.debug(f"Query: {prompt}")
logger.debug(f"System prompt: {system_prompt}")
logger.debug("Full context:")
if "response_format" in kwargs: if "response_format" in kwargs:
response = await openai_async_client.beta.chat.completions.parse( response = await openai_async_client.beta.chat.completions.parse(
model=model, messages=messages, **kwargs model=model, messages=messages, **kwargs

2
lightrag/prompt.py
View File

@@ -8,7 +8,7 @@ PROMPTS["DEFAULT_RECORD_DELIMITER"] = "##"
PROMPTS["DEFAULT_COMPLETION_DELIMITER"] = "<|COMPLETE|>" PROMPTS["DEFAULT_COMPLETION_DELIMITER"] = "<|COMPLETE|>"
PROMPTS["process_tickers"] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"] PROMPTS["process_tickers"] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
PROMPTS["DEFAULT_ENTITY_TYPES"] = ["organization", "person", "geo", "event"] PROMPTS["DEFAULT_ENTITY_TYPES"] = ["organization", "person", "geo", "event", "category"]
PROMPTS["entity_extraction"] = """-Goal- PROMPTS["entity_extraction"] = """-Goal-
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. 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.