From 8f54b62e86b550511a747e02a49ed9057039f1e8 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Fri, 13 Dec 2024 08:50:32 +0800
Subject: [PATCH] Component debugging funcionality. (#4012)

### What problem does this PR solve?

#3993
### Type of change

- [x] New Feature (non-breaking change which adds functionality)
---
 agent/component/categorize.py |  4 ++++
 agent/component/generate.py   | 11 ++++++-----
 agent/component/keyword.py    |  3 +++
 agent/component/relevant.py   |  2 ++
 agent/component/rewrite.py    |  1 +
 5 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/agent/component/categorize.py b/agent/component/categorize.py
index 7b264f131f..6b51eb20d8 100644
--- a/agent/component/categorize.py
+++ b/agent/component/categorize.py
@@ -87,4 +87,8 @@ def _run(self, history, **kwargs):
 
         return Categorize.be_output(list(self._param.category_description.items())[-1][1]["to"])
 
+    def debug(self, **kwargs):
+        df = self._run([], **kwargs)
+        cpn_id = df.iloc[0, 0]
+        return Categorize.be_output(self._canvas.get_compnent_name(cpn_id))
 
diff --git a/agent/component/generate.py b/agent/component/generate.py
index 0dfe8008f3..60ca740cb7 100644
--- a/agent/component/generate.py
+++ b/agent/component/generate.py
@@ -111,9 +111,9 @@ def set_cite(self, retrieval_res, answer):
 
     def get_input_elements(self):
         if self._param.parameters:
-            return [{"key": "user"}, *self._param.parameters]
+            return [{"key": "user", "name": "User"}, *self._param.parameters]
 
-        return [{"key": "user"}]
+        return [{"key": "user", "name": "User"}]
 
     def _run(self, history, **kwargs):
         chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id)
@@ -220,14 +220,15 @@ def stream_output(self, chat_mdl, prompt, retrieval_res):
 
         self.set_output(Generate.be_output(res))
 
-    def debug(self, history, **kwargs):
+    def debug(self, **kwargs):
         chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id)
         prompt = self._param.prompt
 
         for para in self._param.debug_inputs:
-            kwargs[para["key"]] = para["value"]
+            kwargs[para["key"]] = para.get("value", "")
 
         for n, v in kwargs.items():
             prompt = re.sub(r"\{%s\}" % re.escape(n), str(v).replace("\\", " "), prompt)
 
-        return chat_mdl.chat(prompt, [{"role": "user", "content": kwargs.get("user", "")}], self._param.gen_conf())
+        ans = chat_mdl.chat(prompt, [{"role": "user", "content": kwargs.get("user", "")}], self._param.gen_conf())
+        return pd.DataFrame([ans])
diff --git a/agent/component/keyword.py b/agent/component/keyword.py
index 1e417511ed..3b71441c0d 100644
--- a/agent/component/keyword.py
+++ b/agent/component/keyword.py
@@ -60,3 +60,6 @@ def _run(self, history, **kwargs):
         ans = re.sub(r".*keyword:", "", ans).strip()
         logging.debug(f"ans: {ans}")
         return KeywordExtract.be_output(ans)
+
+    def debug(self, **kwargs):
+        return self._run([], **kwargs)
\ No newline at end of file
diff --git a/agent/component/relevant.py b/agent/component/relevant.py
index 65a8433d13..201506afae 100644
--- a/agent/component/relevant.py
+++ b/agent/component/relevant.py
@@ -78,4 +78,6 @@ def _run(self, history, **kwargs):
             return Relevant.be_output(self._param.no)
         assert False, f"Relevant component got: {ans}"
 
+    def debug(self, **kwargs):
+        return self._run([], **kwargs)
 
diff --git a/agent/component/rewrite.py b/agent/component/rewrite.py
index 31390c6f67..f6bf0de695 100644
--- a/agent/component/rewrite.py
+++ b/agent/component/rewrite.py
@@ -110,3 +110,4 @@ def _run(self, history, **kwargs):
         return RewriteQuestion.be_output(ans)
 
 
+