mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-07-03 07:04:04 +00:00
Refactor LLM cache handling and entity extraction
- Removed custom LLM function in entity extraction - Simplified cache handling logic - Added `force_llm_cache` parameter - Updated cache handling conditions
This commit is contained in:
parent
6c7d7c25d3
commit
b45ae1567c
@ -352,32 +352,6 @@ async def extract_entities(
|
|||||||
input_text: str, history_messages: list[dict[str, str]] = None
|
input_text: str, history_messages: list[dict[str, str]] = None
|
||||||
) -> str:
|
) -> str:
|
||||||
if enable_llm_cache_for_entity_extract and llm_response_cache:
|
if enable_llm_cache_for_entity_extract and llm_response_cache:
|
||||||
custom_llm = None
|
|
||||||
if (
|
|
||||||
global_config["embedding_cache_config"]
|
|
||||||
and global_config["embedding_cache_config"]["enabled"]
|
|
||||||
):
|
|
||||||
new_config = global_config.copy()
|
|
||||||
new_config["embedding_cache_config"] = None
|
|
||||||
new_config["enable_llm_cache"] = True
|
|
||||||
|
|
||||||
# create a llm function with new_config for handle_cache
|
|
||||||
async def custom_llm(
|
|
||||||
prompt,
|
|
||||||
system_prompt=None,
|
|
||||||
history_messages=[],
|
|
||||||
keyword_extraction=False,
|
|
||||||
**kwargs,
|
|
||||||
) -> str:
|
|
||||||
# 合并 new_config 和其他 kwargs,保证其他参数不被覆盖
|
|
||||||
merged_config = {**kwargs, **new_config}
|
|
||||||
return await use_llm_func(
|
|
||||||
prompt,
|
|
||||||
system_prompt=system_prompt,
|
|
||||||
history_messages=history_messages,
|
|
||||||
keyword_extraction=keyword_extraction,
|
|
||||||
**merged_config,
|
|
||||||
)
|
|
||||||
|
|
||||||
if history_messages:
|
if history_messages:
|
||||||
history = json.dumps(history_messages, ensure_ascii=False)
|
history = json.dumps(history_messages, ensure_ascii=False)
|
||||||
@ -392,7 +366,7 @@ async def extract_entities(
|
|||||||
_prompt,
|
_prompt,
|
||||||
"default",
|
"default",
|
||||||
cache_type="extract",
|
cache_type="extract",
|
||||||
llm=custom_llm,
|
force_llm_cache=True,
|
||||||
)
|
)
|
||||||
if cached_return:
|
if cached_return:
|
||||||
logger.debug(f"Found cache for {arg_hash}")
|
logger.debug(f"Found cache for {arg_hash}")
|
||||||
|
@ -484,10 +484,10 @@ def dequantize_embedding(
|
|||||||
|
|
||||||
|
|
||||||
async def handle_cache(
|
async def handle_cache(
|
||||||
hashing_kv, args_hash, prompt, mode="default", cache_type=None, llm=None
|
hashing_kv, args_hash, prompt, mode="default", cache_type=None, force_llm_cache=False
|
||||||
):
|
):
|
||||||
"""Generic cache handling function"""
|
"""Generic cache handling function"""
|
||||||
if hashing_kv is None or not hashing_kv.global_config.get("enable_llm_cache"):
|
if hashing_kv is None or not (force_llm_cache or hashing_kv.global_config.get("enable_llm_cache")):
|
||||||
return None, None, None, None
|
return None, None, None, None
|
||||||
|
|
||||||
if mode != "default":
|
if mode != "default":
|
||||||
@ -513,9 +513,7 @@ async def handle_cache(
|
|||||||
similarity_threshold=embedding_cache_config["similarity_threshold"],
|
similarity_threshold=embedding_cache_config["similarity_threshold"],
|
||||||
mode=mode,
|
mode=mode,
|
||||||
use_llm_check=use_llm_check,
|
use_llm_check=use_llm_check,
|
||||||
llm_func=llm
|
llm_func=llm_model_func if use_llm_check else None,
|
||||||
if (use_llm_check and llm is not None)
|
|
||||||
else (llm_model_func if use_llm_check else None),
|
|
||||||
original_prompt=prompt if use_llm_check else None,
|
original_prompt=prompt if use_llm_check else None,
|
||||||
cache_type=cache_type,
|
cache_type=cache_type,
|
||||||
)
|
)
|
||||||
|
Loading…
x
Reference in New Issue
Block a user