From c9481c81b9ebaab16be423e48eec023e0f05ae6f Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Sat, 1 Feb 2025 23:05:02 +0800
Subject: [PATCH] Add cache type "extract" for entity extraction

---
 lightrag/operate.py | 4 ++--
 lightrag/utils.py   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/lightrag/operate.py b/lightrag/operate.py
index 49f6bf39..bc011cb9 100644
--- a/lightrag/operate.py
+++ b/lightrag/operate.py
@@ -391,7 +391,7 @@ async def extract_entities(
                 arg_hash,
                 _prompt,
                 "default",
-                cache_type="default",
+                cache_type="extract",
                 llm=custom_llm,
             )
             if cached_return:
@@ -407,7 +407,7 @@ async def extract_entities(
                 res: str = await use_llm_func(input_text)
             await save_to_cache(
                 llm_response_cache,
-                CacheData(args_hash=arg_hash, content=res, prompt=_prompt),
+                CacheData(args_hash=arg_hash, content=res, prompt=_prompt, cache_type="extract"),
             )
             return res
 
diff --git a/lightrag/utils.py b/lightrag/utils.py
index be6bdafe..d4d42b40 100644
--- a/lightrag/utils.py
+++ b/lightrag/utils.py
@@ -105,7 +105,7 @@ def compute_args_hash(*args, cache_type: str = None) -> str:
     """Compute a hash for the given arguments.
     Args:
         *args: Arguments to hash
-        cache_type: Type of cache (e.g., 'keywords', 'query')
+        cache_type: Type of cache (e.g., 'keywords', 'query', 'extract')
     Returns:
         str: Hash string
     """