from langchain.agents import create_agent
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_core.tools import tool


@tool(description="查询股票价格")
def get_price(name: str) -> str:
    return f"股票{name}的价格是20元"


@tool(description="查询股票信息")
def get_info(name: str) -> str:
    return f"股票{name}是一家A股上市公司"


agent = create_agent(
    model=ChatTongyi(model="qwen3-max"),
    tools=[get_price, get_info],
    system_prompt="你是一个智能助手，可以回答股票相关问题，请告知我思考过程，让我知道你为什么调用某个工具"
)

res = agent.stream(
    {"messages": [{
        "role": "user", "content": "kk公司股价多少，并介绍一下"
    }]},
    stream_mode="values"
)

for chunk in res:
    last_msg = chunk["messages"][-1]
    if last_msg.content:
        print(type(last_msg).__name__, last_msg.content)

    if type(last_msg).__name__ == "AIMessage":
        if last_msg.tool_calls:
            tool_list = [tc['name'] for tc in last_msg.tool_calls]
            print(f"工具调佣 {tool_list}")