Fix concurrent problem on extract_entities function.

- Abandon the approach of temporarily replacing the global llm_model_func configuration
- Introduce custom_llm function with new_config for handle_cache while extracting entities
- Update handle_cache to accept custom_llm

lightrag/operate.py

@@ -352,7 +352,7 @@ async def extract_entities(
         input_text: str, history_messages: list[dict[str, str]] = None
     ) -> str:
         if enable_llm_cache_for_entity_extract and llm_response_cache:
-            need_to_restore = False
+            custom_llm = None            
             if (
                 global_config["embedding_cache_config"]
                 and global_config["embedding_cache_config"]["enabled"]
@@ -360,8 +360,21 @@ async def extract_entities(
                 new_config = global_config.copy()
                 new_config["embedding_cache_config"] = None
                 new_config["enable_llm_cache"] = True
-                llm_response_cache.global_config = new_config
-                need_to_restore = 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:
                 history = json.dumps(history_messages, ensure_ascii=False)
                 _prompt = history + "\n" + input_text
@@ -370,10 +383,13 @@ async def extract_entities(
 
             arg_hash = compute_args_hash(_prompt)
             cached_return, _1, _2, _3 = await handle_cache(
-                llm_response_cache, arg_hash, _prompt, "default", cache_type="default"
+                llm_response_cache, 
+                arg_hash, 
+                _prompt, 
+                "default", 
+                cache_type="default",
+                llm=custom_llm
             )
-            if need_to_restore:
-                llm_response_cache.global_config = global_config
             if cached_return:
                 logger.debug(f"Found cache for {arg_hash}")
                 statistic_data["llm_cache"] += 1