diff --git a/agent/component/rewrite.py b/agent/component/rewrite.py index e0302ec95..0ab19d412 100644 --- a/agent/component/rewrite.py +++ b/agent/component/rewrite.py @@ -14,9 +14,8 @@ # limitations under the License. # from abc import ABC -from api.db import LLMType -from api.db.services.llm_service import LLMBundle from agent.component import GenerateParam, Generate +from rag.prompts import full_question class RewriteQuestionParam(GenerateParam): @@ -33,48 +32,6 @@ class RewriteQuestionParam(GenerateParam): def check(self): super().check() - def get_prompt(self, conv, language, query): - prompt = """ -Role: A helpful assistant -Task: Generate a full user question that would follow the conversation. -Requirements & Restrictions: - - Text generated MUST be in the same language of the original user's question. - - If the user's latest question is completely, don't do anything, just return the original question. - - DON'T generate anything except a refined question.""" - - if language: - prompt += f""" - - Text generated MUST be in {language}""" - - prompt += f""" -###################### --Examples- -###################### -# Example 1 -## Conversation -USER: What is the name of Donald Trump's father? -ASSISTANT: Fred Trump. -USER: And his mother? -############### -Output: What's the name of Donald Trump's mother? ------------- -# Example 2 -## Conversation -USER: What is the name of Donald Trump's father? -ASSISTANT: Fred Trump. -USER: And his mother? -ASSISTANT: Mary Trump. -USER: What's her full name? -############### -Output: What's the full name of Donald Trump's mother Mary Trump? -###################### -# Real Data -## Conversation -{conv} -############### -""" - return prompt - class RewriteQuestion(Generate, ABC): component_name = "RewriteQuestion" @@ -83,15 +40,10 @@ class RewriteQuestion(Generate, ABC): hist = self._canvas.get_history(self._param.message_history_window_size) query = self.get_input() query = str(query["content"][0]) if "content" in query else "" - conv = [] - for m in hist: - if m["role"] not in ["user", "assistant"]: - continue - conv.append("{}: {}".format(m["role"].upper(), m["content"])) - conv = "\n".join(conv) - chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) - ans = chat_mdl.chat(self._param.get_prompt(conv, self.gen_lang(self._param.language), query), - [{"role": "user", "content": "Output: "}], self._param.gen_conf()) + messages = [h for h in hist if h["role"]!="system"] + if messages[-1]["role"] != "user": + messages.append({"role": "user", "content": query}) + ans = full_question(self._canvas.get_tenant_id(), self._param.llm_id, messages, self.gen_lang(self._param.language)) self._canvas.history.pop() self._canvas.history.append(("user", ans)) return RewriteQuestion.be_output(ans) diff --git a/rag/prompts.py b/rag/prompts.py index 4a07f175c..7c0f23be9 100644 --- a/rag/prompts.py +++ b/rag/prompts.py @@ -182,7 +182,7 @@ Requirements: return kwd -def full_question(tenant_id, llm_id, messages): +def full_question(tenant_id, llm_id, messages, language=None): if llm_id2llm_type(llm_id) == "image2text": chat_mdl = LLMBundle(tenant_id, LLMType.IMAGE2TEXT, llm_id) else: @@ -204,9 +204,16 @@ Task and steps: 2. If the user's question involves relative date, you need to convert it into absolute date based on the current date, which is {today}. For example: 'yesterday' would be converted to {yesterday}. Requirements & Restrictions: - - Text generated MUST be in the same language of the original user's question. - If the user's latest question is completely, don't do anything, just return the original question. - - DON'T generate anything except a refined question. + - DON'T generate anything except a refined question.""" + if language: + prompt += f""" + - Text generated MUST be in {language}.""" + else: + prompt += """ + - Text generated MUST be in the same language of the original user's question. +""" + prompt += f""" ###################### -Examples- @@ -239,8 +246,8 @@ ASSISTANT: Cloudy. USER: What's about tomorrow in Rochester? ############### Output: What's the weather in Rochester on {tomorrow}? -###################### +###################### # Real Data ## Conversation {conv}