from langchain_openai import ChatOpenAI
from langchain_core.tools import Tool
from langchain.agents import create_react_agent  # 改用 ReAct 智能体
from langchain.agents import AgentExecutor
from langchain_core.prompts import PromptTemplate  # ReAct 用 PromptTemplate 而非 ChatPromptTemplate
from langchain_core.messages import AIMessage, HumanMessage
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

# 1. 初始化 LLM（保持不变）
DEEPSEEK_API_KEY = "123"  # 替换为实际的 API Key
llm = ChatOpenAI(
    api_key=DEEPSEEK_API_KEY,
    base_url="http://172.25.133.51:8085/v1",
    model="qwen3.5-27b-awq",
    temperature=0.3,
    max_tokens=1024,
)

# 2. 工具函数
def huawei_mall_search(query: str) -> str:
    """华为商城搜索工具"""
    print(f"[DEBUG] 工具被调用！搜索关键词：{query}")
    search_results = {
        "众测活动": "华为商城众测活动是让用户体验新品并反馈意见的活动。目前有Mate 60系列众测，参与可赢取礼品。",
        "手机": "华为商城最新手机：Mate 60系列、P60系列、nova系列等。",
        "笔记本": "华为MateBook X Pro、MateBook D系列笔记本电脑。",
        "手表": "华为Watch 4、Watch GT系列智能手表。",
        "默认": "请在华为商城官网查看详细信息或联系客服。"
    }
    for keyword in search_results:
        if keyword in query:
            return f"华为商城搜索结果：{search_results[keyword]}"
    return search_results["默认"]

# 3. 创建工具
huawei_tool = Tool(
    name="huawei_mall_search",
    description="查询华为商城相关信息，包括产品、活动、政策等",
    func=huawei_mall_search,
)
tools = [huawei_tool]

# 4. 定义 ReAct 提示词模板（关键修改！）
react_prompt = PromptTemplate.from_template("""
Answer the following questions as best you can. You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Previous conversation history:
{chat_history}

Question: {input}
Thought: {agent_scratchpad}
""")

#上面提示词中{chat_history}是记录历史会话记录的不能少

# 5. 创建 ReAct 智能体（关键修改！）
try:
    agent = create_react_agent(llm=llm, tools=tools, prompt=react_prompt)
    print("Agent 创建成功")
except Exception as e:
    print(f"创建 Agent 失败: {e}")
    exit()

# 6. 创建执行器（保持不变）
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=3,
    handle_parsing_errors=True,
    return_intermediate_steps=True
)

# 7. 测试（保持不变）
print("\n" + "=" * 60)
print("测试 Agent 工具调用")
print("=" * 60)



# 1. 用字典存储每个 session_id 的独立历史
session_histories = {}

# 2. 会话存储函数：为每个 session_id 创建/返回独立的历史
def get_session_history(session_id: str) -> ChatMessageHistory:
    if session_id not in session_histories:
        session_histories[session_id] = ChatMessageHistory()
    return session_histories[session_id]

# 3. 创建带历史的智能体
agent_with_chat_history = RunnableWithMessageHistory(
    agent_executor,
    get_session_history,  # 使用新的会话存储函数
    input_messages_key="input",
    history_messages_key="chat_history",
)


# 4. 测试不同 session_id
chat1 = agent_with_chat_history.invoke(
    {"input": "你好，我是小老虎"},
    config={"configurable": {"session_id": "tiger"}},
)

chat2 = agent_with_chat_history.invoke(
    {"input": "你好，我是小兔子"},
    config={"configurable": {"session_id": "rabbit"}},
)

# 5. 查看特定 session_id 的会话记录
print("tiger 的会话记录：")
for msg in session_histories["tiger"].messages:
    print(f"{msg.type}: {msg.content}")

print("\nrabbit 的会话记录：")
for msg in session_histories["rabbit"].messages:
    print(f"{msg.type}: {msg.content}")