Some enhancements:

- Enable the llm_cache storage to support get_by_mode_and_id, to improve the performance for using real KV server
- Provide an option for the developers to cache the LLM response when extracting entities for a document. Solving the paint point that sometimes the process failed, the processed chunks we need to call LLM again, money and time wasted. With the new option (by default not enabled) enabling, we can cache that result, can significantly save the time and money for beginners.

lightrag/utils.py

@@ -454,7 +454,10 @@ async def handle_cache(hashing_kv, args_hash, prompt, mode="default"):
 
     # For naive mode, only use simple cache matching
     if mode == "naive":
-        mode_cache = await hashing_kv.get_by_id(mode) or {}
+        if exists_func(hashing_kv, "get_by_mode_and_id"):
+            mode_cache = await hashing_kv.get_by_mode_and_id(mode, args_hash) or {}
+        else:
+            mode_cache = await hashing_kv.get_by_id(mode) or {}
         if args_hash in mode_cache:
             return mode_cache[args_hash]["return"], None, None, None
         return None, None, None, None
@@ -488,7 +491,10 @@ async def handle_cache(hashing_kv, args_hash, prompt, mode="default"):
             return best_cached_response, None, None, None
     else:
         # Use regular cache
-        mode_cache = await hashing_kv.get_by_id(mode) or {}
+        if exists_func(hashing_kv, "get_by_mode_and_id"):
+            mode_cache = await hashing_kv.get_by_mode_and_id(mode, args_hash) or {}
+        else:
+            mode_cache = await hashing_kv.get_by_id(mode) or {}
         if args_hash in mode_cache:
             return mode_cache[args_hash]["return"], None, None, None
 
@@ -510,7 +516,13 @@ async def save_to_cache(hashing_kv, cache_data: CacheData):
     if hashing_kv is None or hasattr(cache_data.content, "__aiter__"):
         return
 
-    mode_cache = await hashing_kv.get_by_id(cache_data.mode) or {}
+    if exists_func(hashing_kv, "get_by_mode_and_id"):
+        mode_cache = (
+            await hashing_kv.get_by_mode_and_id(cache_data.mode, cache_data.args_hash)
+            or {}
+        )
+    else:
+        mode_cache = await hashing_kv.get_by_id(cache_data.mode) or {}
 
     mode_cache[cache_data.args_hash] = {
         "return": cache_data.content,
@@ -543,3 +555,15 @@ def safe_unicode_decode(content):
     )
 
     return decoded_content
+
+
+def exists_func(obj, func_name: str) -> bool:
+    """Check if a function exists in an object or not.
+    :param obj:
+    :param func_name:
+    :return: True / False
+    """
+    if callable(getattr(obj, func_name, None)):
+        return True
+    else:
+        return False