from langchain_core.prompts import ChatPromptTemplate, PromptTemplate

def get_prompt_template(prompt_type:str) -> ChatPromptTemplate:
    """返回 RAG 问答链的提示模板。"""

    rewrite_prompt = PromptTemplate(
input_variables=["query"],
template="""你是⼀个查询优化助⼿。请将⽤户的⼝语化问题改写为更适合信息检索的
精确查询。
改写要求：
1. 补充隐含的上下⽂信息
2. 将⼝语化表达转为专业表述
在 RAG 链路中集成
将查询重写作为检索前的预处理步骤：
3. 消除歧义，明确查询意图
4. 保持原意不变，不要添加原问题未提及的内容
5. 直接输出改写后的查询，不要解释
⽤户问题: {query}
改写后的查询:"""
)

    decompose_prompt = PromptTemplate(
input_variables=["question"],
template="""你是⼀个问题分解助⼿。请将⽤户的复杂问题分解为 2-4 个独⽴的⼦问
题，
每个⼦问题应该能独⽴检索和回答。
要求：
1. ⼦问题之间互不依赖，可以并⾏检索
2. ⼦问题覆盖原始问题的所有⽅⾯
3. 每个⼦问题简洁明确
4. 以 JSON 数组格式输出
⽤户问题: {question}
输出格式示例: ["⼦问题1", "⼦问题2", "⼦问题3"]
⼦问题列表:"""
)

    clarify_prompt = PromptTemplate(
input_variables=["query"],
template="""你是⼀个智能客服助⼿。请判断⽤户的提问是否包含⾜够的信息来进⾏准
确检索和回答。
判断标准：
1. 问题是否有明确的主语（指代的对象是否清晰）
2. 问题是否包含必要的上下⽂（时间、地点、产品型号等）
3. 问题是否存在歧义（可能有多种理解⽅式）
请以 JSON 格式返回：
- 如果问题清晰：{{"need_clarify": false, "reason": "问题清晰的原因"}}
- 如果需要澄清：{{"need_clarify": true, "clarification": "向⽤户提出的澄清
问题", "assumption": "如果必须回答时的合理假设"}}
⽤户问题: {query}
结果:"""
)
    if prompt_type == "rewrite":
        return rewrite_prompt
    elif prompt_type == "decompose":
        return decompose_prompt
    elif prompt_type == "clarify":
        return clarify_prompt
    else:
        return (ChatPromptTemplate.from_messages(
                [
                    (
                        "system",
                        "你是专业的汽车知识库助手。只能依据给出的上下文回答；"
                        "上下文没有答案时，请直接说明“抱歉，我无法从知识库中回答这个问题”。",
                    ),
                    ("user", "上下文：\n{context}\n\n用户问题：{question}"),
                ]
            ))