diff --git a/agent/component/keyword.py b/agent/component/keyword.py index 3b71441c0..1e7cbacde 100644 --- a/agent/component/keyword.py +++ b/agent/component/keyword.py @@ -57,9 +57,10 @@ class KeywordExtract(Generate, ABC): ans = chat_mdl.chat(self._param.get_prompt(), [{"role": "user", "content": query}], self._param.gen_conf()) + ans = re.sub(r".*", "", ans, flags=re.DOTALL) 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 + return self._run([], **kwargs)