PROMPT_QUERY_REWRITE = """作为一个对话助手，你的任务是的作用是归纳总结和问题相关的历史对话，并优化问题。你需要做的步骤如下：
一、背景信息提取:
·请先浏览并理解对话历史记录。
·识别出与原问题最相关的部分，包括但不限于关键事件、决策点、已解答的问题及其答案等。
二、问题分析与重构:
·根据上述提取的信息，确定原问题的核心意图。
·如果发现原问题不够明确或存在歧义，请尝试重构问题，使其更加清晰、具体。
·重构时考虑是否有必要引入更多上下文来帮助更准确地定位问题。
三、相关信息整合:
·将从历史对话中获得的相关信息以简洁明了的方式融入到问题重构过程中。
·目标是让最终形成的问题不仅反映了原问题的原始意图，同时也包含了有助于快速找到正确答案的所有必要细节。
四、输出格式:
·最终输出应该包含两部分内容：一是经过优化后的问题表述；二是对之前讨论过的相关要点的简短总结。
·请确保你的回答既直接又富有信息量，同时保持语言的专业性和友好性。
例如
'''
'''
原问题: 介绍下剧情。
重写后的问题: 介绍下剧情
----------------
历史记录: 
'''
Q: 我的电脑昨天还好好的，今天早上突然就开不了机了,按电源按钮一点反应都没有。
A: 可能是电源插座问题，请问插座有问题吗
Q: 电源插座检查过了，也没问题
A: 那可能是电源适配器的问题，您可以尝试换一个看看。
'''
原问题: 如果电源适配器没有问题，还有什么可能的原因？
优化后的问题: 如果电源适配器正常工作，电脑无法开机还可能由哪些硬件故障引起？
相关背景总结: 电脑突然无法开机，已经确认电源线和插座没有问题，初步排除了电源供应方面的问题。
----------------
历史记录: 
'''
Q: 对话背景。
A: 当前对话是关于 Nginx 的介绍和使用等。
'''
原问题: 怎么下载
优化后的问题: Nginx 如何怎么下载？
相关背景总结: Nginx 的介绍和使用
----------------
历史记录: 
'''
Q: 护产假多少天?
A: 护产假的天数根据员工所在的城市而定。请提供您所在的城市，以便我回答您的问题。
'''
原问题: 沈阳
优化后的问题: 沈阳的护产假多少天？
相关背景总结: 护产假的规定
----------------
历史记录: 
'''
Q: 作者是谁？
A: FastGPT 的作者是 labring。
'''
原问题: Tell me about him
优化后的问题: Introduce labring, the author of FastGPT
相关背景总结: FastGPT 的作者是 labring
----------------
历史记录:
'''
Q: 对话背景。
A: 关于 FatGPT 的介绍和使用等问题。
'''
原问题: 你好。
优化后的问题: 你好
相关背景总结: 关于 FatGPT 的介绍和使用等问题
----------------
历史记录:
'''
'''
原问题: 武汉的天气怎么样？
优化后的问题: 武汉的天气怎么样？
相关背景总结: 无
----------------
历史记录:
'''
Q: FastGPT 的优势
A: 1. 开源
   2. 简便
   3. 扩展性强
'''
原问题: 介绍下第2点。
优化后的问题: 介绍下 FastGPT 简便的优势
相关背景总结: FastGPT 的优势是开源，简便和扩展性强
----------------
历史记录:
'''
Q: 什么是 FastGPT？
A: FastGPT 是一个 RAG 平台。
Q: 什么是 Laf？
A: Laf 是一个云函数开发平台。
'''
原问题: 它们有什么关系？
优化后的问题: FastGPT和Laf有什么关系？
相关背景总结: 简单介绍了FastGPT和Laf
----------------
历史记录:
'''
{histories}
'''
原问题: {query}
优化后的问题: 
相关背景总结: 
"""


from pydantic import BaseModel,Field
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
import re
from ..utils.logger import get_logger

class ReWriterModel(BaseModel):
    rewriter: str = Field(
        description="优化后的问题",
    )
    summarize: str = Field(
        description="相关背景总结",
    )


class ReWriteLLM:
    def __init__(self,llm):
        parser = PydanticOutputParser(pydantic_object=ReWriterModel)
        prompt = PromptTemplate(
            template=PROMPT_QUERY_REWRITE + "\n{format_instructions}\n",
            input_variables=["histories","query"],
            partial_variables={"format_instructions": parser.get_format_instructions()},
        )
        self.router = prompt | llm
        self.parser = parser
        self.logger = get_logger(self.__class__.__name__)
    def invoke(self, question, history):
        self.logger.info(f"重写用户查询: {question}")
        try:
            response = self.router.invoke({"query": question, "histories": history})
            result = self.parser.invoke(response)
        except:
            # 定义正则表达式模式

            return response.content
        self.logger.info(f"重写后的问题: {result.rewriter}")
        return result.rewriter

    def extend_history(self, history):
        history_extend = ""
        for msg in history:
            if msg['obj'] == 'Human':
                history_extend += f"Q: {msg['value'][0]['text']['content']}\n"
            elif msg['obj'] == 'AI':
                history_extend += f"A: {msg['value'][0]['text']['content']}\n"
        return history_extend

def new_re_rewriter_llm(llm):
    re_writer = ReWriteLLM(llm)
    return re_writer