Refa: add more logs to KG. (#8889)

### What problem does this PR solve?


### Type of change

- [x] Refactoring

graphrag/general/extractor.py

@@ -55,12 +55,18 @@ class Extractor:
         if response:
             return response
         _, system_msg = message_fit_in([{"role": "system", "content": system}], int(self._llm.max_length * 0.92))
-        response = self._llm.chat(system_msg[0]["content"], hist, conf)
-        response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
-        if response.find("**ERROR**") >= 0:
-            logging.warning(f"Extractor._chat got error. response: {response}")
-            return ""
-        set_llm_cache(self._llm.llm_name, system, response, history, gen_conf)
+        for attempt in range(3):
+            try:
+                response = self._llm.chat(system_msg[0]["content"], hist, conf)
+                response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
+                if response.find("**ERROR**") >= 0:
+                    raise Exception(response)
+                set_llm_cache(self._llm.llm_name, system, response, history, gen_conf)
+            except Exception as e:
+                logging.exception(e)
+                if attempt == 2:
+                    raise
+
         return response
 
     def _entities_and_relations(self, chunk_key: str, records: list, tuple_delimiter: str):

graphrag/general/index.py

@@ -39,6 +39,14 @@ from rag.nlp import rag_tokenizer, search
 from rag.utils.redis_conn import RedisDistributedLock
 
 
+@timeout(30, 2)
+async def _is_strong_enough(chat_model, embedding_model):
+    _ = await trio.to_thread.run_sync(lambda: embedding_model.encode(["Are you strong enough!?"]))
+    res =  await trio.to_thread.run_sync(lambda: chat_model.chat("Nothing special.", [{"role":"user", "content": "Are you strong enough!?"}]))
+    if res.find("**ERROR**") >= 0:
+        raise Exception(res)
+
+
 async def run_graphrag(
     row: dict,
     language,
@@ -48,6 +56,11 @@ async def run_graphrag(
     embedding_model,
     callback,
 ):
+    # Pressure test for GraphRAG task
+    async with trio.open_nursery() as nursery:
+        for _ in range(12):
+            nursery.start_soon(_is_strong_enough, chat_model, embedding_model)
+
     start = trio.current_time()
     tenant_id, kb_id, doc_id = row["tenant_id"], str(row["kb_id"]), row["doc_id"]
     chunks = []
@@ -65,7 +78,7 @@ async def run_graphrag(
         doc_id,
         chunks,
         language,
-        row["kb_parser_config"]["graphrag"]["entity_types"],
+        row["kb_parser_config"]["graphrag"].get("entity_types", []),
         chat_model,
         embedding_model,
         callback,