### 定义状态图
 
from langgraph.graph import StateGraph, MessagesState
 
graph_builder = StateGraph(MessagesState)
 
### 定义模型和 chatbot 节点
 
from langchain_openai import ChatOpenAI
 
llm = ChatOpenAI()
 
def chatbot(state: MessagesState):
    return {"messages": [llm.invoke(state["messages"])]}
 
### 构建和编译图
 
from langgraph.graph import END, START
 
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()
 
### 运行
 
from langchain_core.messages import HumanMessage
 
response = graph.invoke(
    {"messages": [HumanMessage(content="合肥今天天气怎么样？")]}
)
response["messages"][-1].pretty_print()

graph_builder = StateGraph(MessagesState)

class MessagesState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

from typing import TypedDict
 
class State(TypedDict):
    foo: int
    bar: list[str]

from typing import TypedDict, Annotated
from operator import add
 
class State(TypedDict):
    foo: int
    bar: Annotated[list[str], add]

def chatbot(state: MessagesState):
    return {"messages": [llm.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)

graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

for event in graph.stream({"messages": ("user", "合肥今天天气怎么样？")}):
    for value in event.values():
        value["messages"][-1].pretty_print()

### 定义工具
 
from pydantic import BaseModel, Field
from langchain_core.tools import tool
 
class GetWeatherSchema(BaseModel):
    city: str = Field(description = "城市名称，如合肥、北京、上海等")
    date: str = Field(description = "日期，如今天、明天等")
 
@tool(args_schema = GetWeatherSchema)
def get_weather(city: str, date: str):
    """查询天气"""
    if city == "合肥":
        return "今天晴天，气温30度。"
    return "今天有小雨，气温25度。"

### 定义状态图
 
from langgraph.graph import StateGraph, MessagesState
 
graph_builder = StateGraph(MessagesState)

### 定义 tools 节点
 
from langgraph.prebuilt import ToolNode
 
tools = [get_weather]
tool_node = ToolNode(tools)
 
### 定义模型和 chatbot 节点
 
from langchain_openai import ChatOpenAI
 
llm = ChatOpenAI()
llm = llm.bind_tools(tools)
 
def chatbot(state: MessagesState):
    return {"messages": [llm.invoke(state["messages"])]}

### 构建和编译图
 
from langgraph.graph import END, START
from langgraph.prebuilt import tools_condition
 
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("tools", 'chatbot')
graph_builder.add_conditional_edges("chatbot", tools_condition)
graph = graph_builder.compile()

def tools_condition(
    state: Union[list[AnyMessage], dict[str, Any]],
) -> Literal["tools", "__end__"]:
 
    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
    return "__end__"

tools_by_name = {tool.name: tool for tool in tools}
def tool_node(state: dict):
    result = []
    for tool_call in state["messages"][-1].tool_calls:
        tool = tools_by_name[tool_call["function"]["name"]]
        observation = tool.invoke(tool_call["function"]["arguments"])
        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
    return {"messages": result}

### 定义模型和 chatbot 节点
 
from langchain_openai import ChatOpenAI
 
llm = ChatOpenAI(model="gpt-4")
llm = llm.bind_functions(tools)
 
def chatbot(state: MessagesState):
    return {"messages": [llm.invoke(state["messages"])]}

def tools_condition(
    state: MessagesState,
) -> Literal["tools", "__end__"]:
 
    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    if "function_call" in ai_message.additional_kwargs:
        return "tools"
    return "__end__"

import json
from langchain_core.messages import FunctionMessage
 
class BasicToolNode:
 
    def __init__(self, tools: list) -> None:
        self.tools_by_name = {tool.name: tool for tool in tools}
 
    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        outputs = []
        if "function_call" in message.additional_kwargs:
            tool_call = message.additional_kwargs["function_call"]
            tool_result = self.tools_by_name[tool_call["name"]].invoke(
                json.loads(tool_call["arguments"])
            )
            outputs.append(
                FunctionMessage(
                    content=json.dumps(tool_result),
                    name=tool_call["name"]
                )
            )
        return {"messages": outputs}
 
tools = [get_weather]
tool_node = BasicToolNode(tools=tools)

from langgraph.checkpoint.memory import MemorySaver
 
memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "1"}}
 
for event in graph.stream({"messages": ("user", "合肥今天天气怎么样？")}, config):
    for value in event.values():
        value["messages"][-1].pretty_print()
 
for event in graph.stream({"messages": ("user", "要带伞吗？")}, config):
    for value in event.values():
        value["messages"][-1].pretty_print()

from langgraph.checkpoint.postgres import PostgresSaver
 
DB_URI = "postgresql://postgres:123456@localhost:5432/postgres?sslmode=disable"
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
     
    # 第一次运行时初始化
    checkpointer.setup()
     
    graph = graph_builder.compile(checkpointer=checkpointer)
    config = {"configurable": {"thread_id": "1"}}
    for event in graph.stream({"messages": ("user", "合肥今天天气怎么样？")}, config):
        for value in event.values():
            value["messages"][-1].pretty_print()
    for event in graph.stream({"messages": ("user", "要带伞吗？")}, config):
        for value in event.values():
            value["messages"][-1].pretty_print()

from psycopg import Connection
 
connection_kwargs = {
    "autocommit": True,
    "prepare_threshold": 0,
}
with Connection.connect(DB_URI, **connection_kwargs) as conn:
    checkpointer = PostgresSaver(conn)
    graph = graph_builder.compile(checkpointer=checkpointer)
    ...

from psycopg_pool import ConnectionPool
 
with ConnectionPool(conninfo=DB_URI, max_size=20, kwargs=connection_kwargs) as pool:
    checkpointer = PostgresSaver(pool)
    graph = graph_builder.compile(checkpointer=checkpointer)
    ...

class Wrapper:
    ''' 包装类，用于调试 OpenAI 接口的原始入参和出参
    '''
    def __init__(self, wrapped_class):
        self.wrapped_class = wrapped_class
 
    def __getattr__(self, attr):
        original_func = getattr(self.wrapped_class, attr)
 
        def wrapper(*args, **kwargs):
            print(f"Calling function: {attr}")
            print(f"Arguments: {args}, {kwargs}")
            result = original_func(*args, **kwargs)
            print(f"Response: {result}")
            return result
        return wrapper
 
from langchain_openai import ChatOpenAI
 
llm = ChatOpenAI(model="gpt-4")
llm.client = Wrapper(llm.client)
llm = llm.bind_functions(tools)

export LANGCHAIN_TRACING_V2=true
export LANGCHAIN_API_KEY=lsv2_pt_xxx
export LANGCHAIN_ENDPOINT=https://api.smith.langchain.com
export LANGCHAIN_PROJECT=default

class BookTicketSchema(BaseModel):
    from_city: str = Field(description = "出发城市名称，如合肥、北京、上海等")
    to_city: str = Field(description = "到达城市名称，如合肥、北京、上海等")
    date: str = Field(description = "日期，如今天、明天等")
 
@tool(args_schema = BookTicketSchema)
def book_ticket(from_city: str, to_city: str, date: str):
    """预定机票"""
    return "您已成功预定 %s 从 %s 到 %s 的机票" % (date, from_city, to_city)

for event in graph.stream({"messages": ("user", "帮我预定一张明天从合肥到北京的机票")}, config):
    for value in event.values():
        value["messages"][-1].pretty_print()

graph = graph_builder.compile(
    checkpointer=memory,
    interrupt_before=["tools"]
)

snapshot = graph.get_state(config)
print(snapshot.values["messages"][-1])
print(snapshot.next)

### 继续运行
 
for event in graph.stream(None, config):
    for value in event.values():
        value["messages"][-1].pretty_print()

### 继续运行
 
for event in graph.stream(None, config):
    for value in event.values():
        value["messages"][-1].pretty_print()

snapshot = graph.get_state(config)
existing_message = snapshot.values["messages"][-1]
new_messages = [
    # The LLM API expects some ToolMessage to match its tool call. We'll satisfy that here.
    ToolMessage(content="预定失败", tool_call_id=existing_message.tool_calls[0]["id"]),
    # And then directly "put words in the LLM's mouth" by populating its response.
    AIMessage(content="预定失败"),
]
graph.update_state(config, {"messages": new_messages})