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

Fix linting

Browse Source
This commit is contained in:
yangdx
2025-02-02 01:56:32 +08:00
parent b45ae1567c
commit 5d14ab03eb
2 changed files with 22 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
lightrag/operate.py
View File

@@ -352,7 +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:
if history_messages: if history_messages:
history = json.dumps(history_messages, ensure_ascii=False) history = json.dumps(history_messages, ensure_ascii=False)
_prompt = history + "\n" + input_text _prompt = history + "\n" + input_text
@@ -381,7 +380,12 @@ async def extract_entities(
res: str = await use_llm_func(input_text) res: str = await use_llm_func(input_text)
await save_to_cache( await save_to_cache(
llm_response_cache, llm_response_cache,
CacheData(args_hash=arg_hash, content=res, prompt=_prompt, cache_type="extract"), CacheData(
args_hash=arg_hash,
content=res,
prompt=_prompt,
cache_type="extract",
),
) )
return res return res
@@ -747,7 +751,10 @@ async def extract_keywords_only(
# 7. Cache only the processed keywords with cache type # 7. Cache only the processed keywords with cache type
if hl_keywords or ll_keywords: if hl_keywords or ll_keywords:
cache_data = {"high_level_keywords": hl_keywords, "low_level_keywords": ll_keywords} cache_data = {
"high_level_keywords": hl_keywords,
"low_level_keywords": ll_keywords,
}
await save_to_cache( await save_to_cache(
hashing_kv, hashing_kv,
CacheData( CacheData(

15
lightrag/utils.py
View File

@@ -484,10 +484,17 @@ def dequantize_embedding(
async def handle_cache( async def handle_cache(
hashing_kv, args_hash, prompt, mode="default", cache_type=None, force_llm_cache=False 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 (force_llm_cache or 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":
@@ -504,7 +511,9 @@ async def handle_cache(
# Use embedding cache # Use embedding cache
current_embedding = await hashing_kv.embedding_func([prompt]) current_embedding = await hashing_kv.embedding_func([prompt])
llm_model_func = ( llm_model_func = (
hashing_kv.llm_model_func if hasattr(hashing_kv, "llm_model_func") else None hashing_kv.llm_model_func
if hasattr(hashing_kv, "llm_model_func")
else None
) )
quantized, min_val, max_val = quantize_embedding(current_embedding[0]) quantized, min_val, max_val = quantize_embedding(current_embedding[0])
best_cached_response = await get_best_cached_response( best_cached_response = await get_best_cached_response(