From 73a03287a57a1fafdfbddbe65e05a1d5c3f3d807 Mon Sep 17 00:00:00 2001
From: H <43509927+guoyuhao2330@users.noreply.github.com>
Date: Mon, 19 Aug 2024 20:01:41 +0800
Subject: [PATCH] Fix mutiple retrieval component content (#2006)

### What problem does this PR solve?

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 agent/component/base.py     |  2 --
 agent/component/generate.py | 13 +++++++++----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/agent/component/base.py b/agent/component/base.py
index 88c85c3f..10332501 100644
--- a/agent/component/base.py
+++ b/agent/component/base.py
@@ -448,8 +448,6 @@ class ComponentBase(ABC):
             if self.component_name.lower() == "generate" and self.get_component_name(u) == "retrieval":
                 o = self._canvas.get_component(u)["obj"].output(allow_partial=False)[1]
                 if o is not None:
-                    if not "".join(o["content"]):
-                        continue
                     upstream_outs.append(o)
                     continue
             if u not in self._canvas.get_component(self._id)["upstream"]: continue
diff --git a/agent/component/generate.py b/agent/component/generate.py
index 25252652..dab99632 100644
--- a/agent/component/generate.py
+++ b/agent/component/generate.py
@@ -66,6 +66,8 @@ class Generate(ComponentBase):
         return cpnts
 
     def set_cite(self, retrieval_res, answer):
+        retrieval_res.dropna(subset=["vector", "content_ltks"], inplace=True)
+        retrieval_res["empty_response"].fillna("", inplace=True)
         answer, idx = retrievaler.insert_citations(answer, [ck["content_ltks"] for _, ck in retrieval_res.iterrows()],
                                                    [ck["vector"] for _, ck in retrieval_res.iterrows()],
                                                    LLMBundle(self._canvas.get_tenant_id(), LLMType.EMBEDDING,
@@ -117,8 +119,10 @@ class Generate(ComponentBase):
             "obj"].component_name.lower() == "answer":
             return partial(self.stream_output, chat_mdl, prompt, retrieval_res)
 
-        if "empty_response" in retrieval_res.columns:
-            return Generate.be_output(input)
+        if "empty_response" in retrieval_res.columns and not "".join(retrieval_res["content"]):
+            res = {"content": "\n- ".join(retrieval_res["empty_response"]) if "\n- ".join(
+                retrieval_res["empty_response"]) else "Nothing found in knowledgebase!", "reference": []}
+            return Generate.be_output(res)
 
         ans = chat_mdl.chat(prompt, self._canvas.get_history(self._param.message_history_window_size),
                             self._param.gen_conf())
@@ -130,8 +134,9 @@ class Generate(ComponentBase):
 
     def stream_output(self, chat_mdl, prompt, retrieval_res):
         res = None
-        if "empty_response" in retrieval_res.columns and not "\n- ".join(retrieval_res["content"]):
-            res = {"content": "\n- ".join(retrieval_res["empty_response"]), "reference": []}
+        if "empty_response" in retrieval_res.columns and not "".join(retrieval_res["content"]):
+            res = {"content": "\n- ".join(retrieval_res["empty_response"]) if "\n- ".join(
+                retrieval_res["empty_response"]) else "Nothing found in knowledgebase!", "reference": []}
             yield res
             self.set_output(res)
             return