from typing import Annotated, Literal, TypedDict

from langchain_core.messages import HumanMessage
from langchain_anthropic import ChatAnthropic
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode


# 定义代理要使用的工具
@tool
def search(query: str):
    """调用搜索工具."""
    # 这是一个占位符，但不要告诉 LLM...
    if "sf" in query.lower() or "san francisco" in query.lower():
        return "温度为 60 度，有雾。"
    return "温度为 90 度，阳光明媚。"


tools = [search]

tool_node = ToolNode(tools)

model = ChatAnthropic(model="claude-3-5-sonnet-20240620", temperature=0).bind_tools(tools)

# 定义函数来确定是否继续
def should_continue(state: MessagesState) -> Literal["tools", END]:
    messages = state['messages']
    last_message = messages[-1]
    # 如果 LLM 进行工具调用，则路由到 "tools" 节点
    if last_message.tool_calls:
        return "tools"
    # 否则，我们停止（回复用户）
    return END


# 定义调用模型的函数
def call_model(state: MessagesState):
    messages = state['messages']
    response = model.invoke(messages)
    # 我们返回一个列表，因为这将被添加到现有列表中
    return {"messages": [response]}


# 定义一个新的图
workflow = StateGraph(MessagesState)

# 定义我们将循环遍历的两个节点
workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

# 将入口点设置为 `agent`
# 这意味着此节点是第一个被调用的节点
workflow.add_edge(START, "agent")

# 我们现在添加一个条件边
workflow.add_conditional_edges(
    # 首先，我们定义起始节点。我们使用 `agent`。
    # 这意味着这些是在调用 `agent` 节点后采取的边。
    "agent",
    # 接下来，我们传入一个函数，该函数将确定下一个被调用的节点。
    should_continue,
)

# 我们现在从 `tools` 添加一条普通边到 `agent`。
# 这意味着在调用 `tools` 后，下一个被调用的节点是 `agent` 节点。
workflow.add_edge("tools", 'agent')

# 初始化内存以在图运行之间持久化状态
checkpointer = MemorySaver()

# 最后，我们编译它！
# 这将其编译成一个 LangChain Runnable，
# 意味着您可以像使用任何其他可运行程序一样使用它。
# 请注意，我们（可选地）在编译图时传递内存
app = workflow.compile(checkpointer=checkpointer)

# 使用 Runnable
final_state = app.invoke(
    {"messages": [HumanMessage(content="旧金山的 weather 情况如何？")]},
    config={"configurable": {"thread_id": 42}}
)
final_state["messages"][-1].content

final_state = app.invoke(
    {"messages": [HumanMessage(content="纽约怎么样？")]},
    config={"configurable": {"thread_id": 42}}
)
final_state["messages"][-1].content