Fix JSON parsing error

HKUDS · Dec 5, 2024 · 5e1f317 · 5e1f317
1 parent 68fed9b
commit 5e1f317
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/examples/lightrag_lmdeploy_demo.py b/examples/lightrag_lmdeploy_demo.py
@@ -12,7 +12,11 @@
 
 
 async def lmdeploy_model_complete(
-    prompt=None, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+    prompt=None,
+    system_prompt=None,
+    history_messages=[],
+    keyword_extraction=False,
+    **kwargs,
 ) -> str:
     model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
     return await lmdeploy_model_if_cache(

diff --git a/lightrag/operate.py b/lightrag/operate.py
@@ -479,7 +479,7 @@ async def kg_query(
     print(result)
     try:
         # json_text = locate_json_string_body_from_string(result) # handled in use_model_func
-        result = re.search(r"{.*}", json_text, re.DOTALL)
+        result = re.search(r"{.*}", result, re.DOTALL)
         keywords_data = json.loads(result)
         hl_keywords = keywords_data.get("high_level_keywords", [])
         ll_keywords = keywords_data.get("low_level_keywords", [])