import os
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.tools import tool
from langchain_core.prompts import PromptTemplate

# 设置 API Key（也可以通过环境变量设置）
os.environ["OPENAI_API_KEY"] = "your-api-key-here"

# ========== 定义工具 ==========

@tool
def get_weather(city: str) -> str:
    """查询指定城市的当前天气信息。输入城市名称，返回天气描述。"""
    # 这里用模拟数据，实际项目中可以接入真实天气 API
    weather_data = {
        "北京": "晴天，最高气温32℃，最低气温22℃，东南风3级",
        "上海": "多云，最高气温28℃，最低气温20℃，微风",
        "深圳": "阵雨，最高气温30℃，最低气温25℃，南风2级",
    }
    return weather_data.get(city, f"未找到{city}的天气信息")

@tool
def calculator(expression: str) -> str:
    """计算数学表达式。输入一个数学表达式字符串，返回计算结果。"""
    try:
        result = eval(expression)  # 注意：生产环境请用更安全的计算库
        return str(result)
    except Exception as e:
        return f"计算出错：{str(e)}"

# ========== 设置 LLM ==========

llm = ChatOpenAI(model="gpt-4", temperature=0)

# ========== ReAct Prompt 模板 ==========

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!

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

# ========== 创建 ReAct Agent ==========

tools = [get_weather, calculator]
agent = create_react_agent(llm, tools, react_prompt)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,   # 打印每一步的思考过程
    max_iterations=10  # 防止无限循环
)

# ========== 运行示例 ==========

if __name__ == "__main__":
    # 示例1：单步查询
    print("=" * 50)
    result = agent_executor.invoke({"input": "北京今天天气怎么样？"})
    print(f"最终回答：{result['output']}")

    # 示例2：多步推理（需要先查天气，再计算）
    print("=" * 50)
    result = agent_executor.invoke({
        "input": "北京今天最高气温是多少？如果每天省下3度电，一周能省多少度？"
    })
    print(f"最终回答：{result['output']}")