From 4ba4f622a573fe1df2b416635b3dda0f662eed85 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Tue, 31 Dec 2024 14:31:31 +0800
Subject: [PATCH] Refactor (#4303)

### What problem does this PR solve?

### Type of change
- [x] Refactoring
---
 agent/component/base.py             | 2 +-
 api/apps/__init__.py                | 4 ++--
 api/apps/conversation_app.py        | 4 ----
 api/db/services/document_service.py | 7 ++-----
 api/utils/api_utils.py              | 1 +
 rag/svr/task_executor.py            | 8 +++++---
 6 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/agent/component/base.py b/agent/component/base.py
index 8ee7b8186..fa9ddc6a4 100644
--- a/agent/component/base.py
+++ b/agent/component/base.py
@@ -457,7 +457,7 @@ class ComponentBase(ABC):
 
     def get_input(self):
         if self._param.debug_inputs:
-            return pd.DataFrame([{"content": v["value"]} for v in self._param.debug_inputs])
+            return pd.DataFrame([{"content": v["value"]} for v in self._param.debug_inputs if v.get("value")])
 
         reversed_cpnts = []
         if len(self._canvas.path) > 1:
diff --git a/api/apps/__init__.py b/api/apps/__init__.py
index e2916f1c8..42fc5d616 100644
--- a/api/apps/__init__.py
+++ b/api/apps/__init__.py
@@ -152,8 +152,8 @@ def load_user(web_request):
                 return user[0]
             else:
                 return None
-        except Exception:
-            logging.exception("load_user got exception")
+        except Exception as e:
+            logging.warning(f"load_user got exception {e}")
             return None
     else:
         return None
diff --git a/api/apps/conversation_app.py b/api/apps/conversation_app.py
index 222a3179b..c98abd69c 100644
--- a/api/apps/conversation_app.py
+++ b/api/apps/conversation_app.py
@@ -65,10 +65,6 @@ def set_conversation():
             "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
         }
         ConversationService.save(**conv)
-        e, conv = ConversationService.get_by_id(conv["id"])
-        if not e:
-            return get_data_error_result(message="Fail to new a conversation!")
-        conv = conv.to_dict()
         return get_json_result(data=conv)
     except Exception as e:
         return server_error_response(e)
diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py
index ba146fa4f..642ce2425 100644
--- a/api/db/services/document_service.py
+++ b/api/db/services/document_service.py
@@ -96,14 +96,11 @@ class DocumentService(CommonService):
     def insert(cls, doc):
         if not cls.save(**doc):
             raise RuntimeError("Database error (Document)!")
-        e, doc = cls.get_by_id(doc["id"])
-        if not e:
-            raise RuntimeError("Database error (Document retrieval)!")
-        e, kb = KnowledgebaseService.get_by_id(doc.kb_id)
+        e, kb = KnowledgebaseService.get_by_id(doc["kb_id"])
         if not KnowledgebaseService.update_by_id(
                 kb.id, {"doc_num": kb.doc_num + 1}):
             raise RuntimeError("Database error (Knowledgebase)!")
-        return doc
+        return Document(**doc)
 
     @classmethod
     @DB.connection_context()
diff --git a/api/utils/api_utils.py b/api/utils/api_utils.py
index 2dba15bbe..f33588584 100644
--- a/api/utils/api_utils.py
+++ b/api/utils/api_utils.py
@@ -98,6 +98,7 @@ def get_exponential_backoff_interval(retries, full_jitter=False):
 
 def get_data_error_result(code=settings.RetCode.DATA_ERROR,
                           message='Sorry! Data missing!'):
+    logging.exception(Exception(message))
     result_dict = {
         "code": code,
         "message": message}
diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py
index c1f7f8de3..28adc628b 100644
--- a/rag/svr/task_executor.py
+++ b/rag/svr/task_executor.py
@@ -92,10 +92,12 @@ DONE_TASKS = 0
 FAILED_TASKS = 0
 CURRENT_TASK = None
 
+
 class TaskCanceledException(Exception):
     def __init__(self, msg):
         self.msg = msg
 
+
 def set_progress(task_id, from_page=0, to_page=-1, prog=None, msg="Processing..."):
     global PAYLOAD
     if prog is not None and prog < 0:
@@ -250,7 +252,7 @@ def build_chunks(task, progress_callback):
             STORAGE_IMPL.put(task["kb_id"], d["id"], output_buffer.getvalue())
             el += timer() - st
         except Exception:
-            logging.exception("Saving image of chunk {}/{}/{} got exception".format(task["location"], task["name"], d["_id"]))
+            logging.exception("Saving image of chunk {}/{}/{} got exception".format(task["location"], task["name"], d["id"]))
             raise
 
         d["img_id"] = "{}-{}".format(task["kb_id"], d["id"])
@@ -312,6 +314,8 @@ def embedding(docs, mdl, parser_config=None, callback=None):
         if not c:
             c = d["content_with_weight"]
         c = re.sub(r"</?(table|td|caption|tr|th)( [^<>]{0,12})?>", " ", c)
+        if not c:
+            c = "None"
         cnts.append(c)
 
     tk_count = 0
@@ -394,8 +398,6 @@ def run_raptor(row, chat_mdl, embd_mdl, callback=None):
     return res, tk_count, vector_size
 
 
-
-
 def do_handle_task(task):
     task_id = task["id"]
     task_from_page = task["from_page"]