hitl_demo

from typing import TypedDict, Annotated, Any, Literal
from langchain_core.runnables import RunnableConfig
import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from pydantic import BaseModel,Field
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

dotenv.load_dotenv()

class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class DallEArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的问题提示（Prompt）")

# 定义工具和工具列表
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本谷歌搜索API"
        "当你需要回答有关时事问题时，可以调用该工具"
        "该工具的输入是搜索查询语句"
    ),
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)

dalle = OpenAIDALLEImageGenerationTool(
    name="openai_dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    args_schema=DallEArgsSchema
)

# define graph state
class State(TypedDict):
    """
    图状态数据结构，类型为字典
    """
    messages: Annotated[list, add_messages]

# define tools
tools = [google_serper, dalle]
# define llm
llm = ChatOpenAI(model="gpt-4o-mini")

# bind llm with tools
llm_with_tools = llm.bind_tools(tools)

# define chatbot node
def chatbot(state: State, config: RunnableConfig) -> Any:
    """
    聊天机器人函数
    """
    # 1. 获取状态中存储的消息列表数据并传入给LLM
    ai_message = llm_with_tools.invoke(state["messages"])
    # 2. 返回更新/生成的状态
    return {
        "messages": [
            ai_message
        ]
    }


# define route node
def route(state: State, config: RunnableConfig) -> Literal["tools", "__end__"]:
    """
    动态选择工具执行或者结果
    """
    # 1. 获取生成的最后一条消息
    ai_message = state["messages"][-1]
    # 2. 检测消息是否存在 tool_calls 参数，若是则执行"工具路由"
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
    # 3. 否则生成的内容是文本消息，则跳转到结束路由
    return END

# 1. 创建状态图，并使用GraphState作为状态数据
graph_builder = StateGraph(State)

"""
★★★ 路由Route不作为节点 只作为单纯的路由进行节点路径选择
"""

# 2. add nodes
graph_builder.add_node("llm", chatbot)
graph_builder.add_node("tools", ToolNode(tools=tools))

# 3. add edges
graph_builder.add_edge(START, "llm")
graph_builder.add_edge("tools", "llm")
graph_builder.add_conditional_edges("llm", route)


# 4. 编译图为Runnable可运行组件  在tools中工具执行之前进行中断处理
checkpointer = MemorySaver()
graph = graph_builder.compile(checkpointer=checkpointer, interrupt_before=["tools"])


# 5. 调用图结构应用
config = {
    "configurable": {
        "thread_id": "user-001"
    }
}

# 6. 调用图架构应用
state = graph.invoke(
    {"messages": [
        ("human", "2025年北京马拉松前3名成绩是多少？？？")
    ]}, config=config
)

print(state)

# 6. 获取人类tips
if hasattr(state["messages"][-1], "tool_calls") and len(state["messages"][-1].tool_calls) > 0:
    tool_calls = state["messages"][-1].tool_calls
    print("------------- 准备调用工具 ------------", tool_calls)
    print(f"工具调用信息: {tool_calls}")
    print(f"工具名称: {tool_calls[0].get('name') if tool_calls else 'N/A'}")
    print(f"工具参数: {tool_calls[0].get('args') if tool_calls else 'N/A'}")
    human_input = input("若需要调用工具请恢复yes，否则恢复no：")
    if human_input.lower() == "yes":
        print("\n继续执行工具调用...")
        # 继续执行图（从断点处恢复，执行 tools 节点）
        final_state = graph.invoke(None, config)
        print("\n最终执行结果:")
        print(final_state["messages"][-1].pretty_print())
    else:
        print("\n用户拒绝了工具调用，程序结束")
else:
    print("没有需要调用的工具，程序结束")