# Agent 框架核心组件
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.llms.openai_like import OpenAILike
from llama_index.core.workflow import Context
from config import envconfig
# 上下文序列化工具，用于持久化对话状态
from llama_index.core.workflow import JsonPickleSerializer, JsonSerializer

# 初始化通义千问 LLM，必须开启 is_function_calling_model 才能使用 Agent
llm = OpenAILike(
    model=envconfig.ALIYUN_LLM_MODEL,
    api_base=envconfig.DASHSCOPE_BASE_URL,
    api_key=envconfig.DASHSCOPE_API_KEY,
    is_chat_model=True,
    is_function_calling_model=True,
)

from llama_index.core.agent.workflow import FunctionAgent
from config import envconfig

from llama_index.llms.openai_like import OpenAILike


llm = OpenAILike(
    model=envconfig.ALIYUN_LLM_MODEL,
    api_base=envconfig.DASHSCOPE_BASE_URL,
    api_key=envconfig.DASHSCOPE_API_KEY,
    is_chat_model=True,
    is_function_calling_model=True,
)

# 工具函数：docstring 会作为工具描述传递给 LLM，帮助它理解何时调用
def multiply(a: float, b: float) -> float:
    """Multiply two numbers and returns the product"""
    return a * b


def add(a: float, b: float) -> float:
    """Add two numbers and returns the sum"""
    return a + b


# 创建 Agent，绑定工具函数和 LLM
workflow = FunctionAgent(
    tools=[multiply, add],
    llm=llm,
    system_prompt="You are an agent that can perform basic mathematical operations using tools.",
)
# Agent 自动推理：先调 multiply(2,4)=8，再调 add(20,8)=28
response = await workflow.run(user_msg="What is 20+(2*4)?")
print(response)

from llama_index.llms.openai_like import OpenAILike

from llama_index.core.agent.workflow import AgentWorkflow
# 雅虎财经工具集，提供股票行情、新闻等金融数据
from llama_index.tools.yahoo_finance import YahooFinanceToolSpec

def multiply(a: float, b: float) -> float:
    """Multiply two numbers and returns the product"""
    return a * b

def add(a: float, b: float) -> float:
    """Add two numbers and returns the sum"""
    return a + b

llm = OpenAILike(
    model=envconfig.ALIYUN_LLM_MODEL,
    api_base=envconfig.DASHSCOPE_BASE_URL,
    api_key=envconfig.DASHSCOPE_API_KEY,
    is_chat_model=True,
    is_function_calling_model=True,
)

# 将雅虎财经 API 转换为 Agent 可用的工具列表
finance_tools = YahooFinanceToolSpec().to_tool_list()
# 合并自定义数学工具
finance_tools.extend([multiply, add])

# 从工具列表自动创建工作流（自动选择合适的工具）
workflow = AgentWorkflow.from_tools_or_functions(
    finance_tools,
    llm=llm,
    system_prompt="You are an agent that can perform basic mathematical operations using tools."
)

response = await workflow.run(user_msg="中国的股票“阿里巴巴（美股）”近三天行情怎么样")
print(response)

from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai_like import OpenAILike

# DuckDuckGo 搜索工具集
tool_spec = DuckDuckGoSearchToolSpec()


llm = OpenAILike(
    model=envconfig.ALIYUN_LLM_MODEL,
    api_base=envconfig.DASHSCOPE_BASE_URL,
    api_key=envconfig.DASHSCOPE_API_KEY,
    is_chat_model=True,
    is_function_calling_model=True,
)

# 方式一：通过 Agent 自动调用搜索工具
agent = FunctionAgent(
    tools=tool_spec.to_tool_list(),
    llm=llm,
)

print(await agent.run("谁是中国首富"))
# print(await agent.run("怎么样才能去微软工作"))

# 方式二：直接调用搜索工具，跳过 Agent 推理
tool_directly_call = DuckDuckGoSearchToolSpec().duckduckgo_instant_search("中国有多少个民族") # duckduckgo_full_search
print(tool_directly_call)

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.workflow import Context
from llama_index.core.workflow import JsonPickleSerializer, JsonSerializer

llm = OpenAILike(
    model=envconfig.ALIYUN_LLM_MODEL,
    api_base=envconfig.DASHSCOPE_BASE_URL,
    api_key=envconfig.DASHSCOPE_API_KEY,
    is_chat_model=True,
    is_function_calling_model=True,
)


workflow = FunctionAgent(
    llm=llm,
    system_prompt="你是一个知情达理的AI助手",
)


# 创建上下文，自动维护对话历史
ctx = Context(workflow)

# 第一轮：告诉 Agent 域名信息，Agent 会记住
response = await workflow.run(user_msg="你好，我的网站域名是heiseven.top",ctx=ctx)
print(response)
print("="*60)

# 第二轮：Agent 基于上下文回答之前提到的信息
response2 = await workflow.run(user_msg="我有什么网站？",ctx=ctx)
print(response2)
print("="*60)

# 将上下文序列化为字典（可保存到文件/数据库）
ctx_dict = ctx.to_dict(serializer=JsonSerializer())

# 从字典恢复上下文（实现会话持久化）
restored_ctx = Context.from_dict(
    workflow, ctx_dict, serializer=JsonSerializer()
)

# 使用恢复的上下文继续对话
response3 = await workflow.run(user_msg="我这个域名怎么样",ctx=restored_ctx)
print(response3)

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai_like import OpenAILike
from llama_index.core.workflow import Context


llm = OpenAILike(
    model=envconfig.ALIYUN_LLM_MODEL,
    api_base=envconfig.DASHSCOPE_BASE_URL,
    api_key=envconfig.DASHSCOPE_API_KEY,
    is_chat_model=True,
    is_function_calling_model=True,
)

"""
当你使用 AgentWorkflow.from_tools_or_functions(...) 或 FunctionAgent 并传入 initial_state={"name": "unset"} 时，框架在底层会自动将你的字典包装在一个名为 "state" 的键下

{
    "state": {
        "name": "unset"
    }
}

"""

# 工具函数：通过 ctx 读写状态
async def set_name(ctx: Context, name: str) -> str:
    state = await ctx.get("state")  # 读取当前状态
    state["name"] = name            # 更新状态
    await ctx.set("state", state)   # 写回状态
    return f"Name set to {name}"

# initial_state 初始化状态，框架会自动包装为 {"state": {...}}
workflow = FunctionAgent(
    tools = [set_name],
    llm=llm,
    system_prompt="You are a helpful assistant that can set a name.",
    initial_state={"name": "unset"},
)

ctx = Context(workflow)

# 查询状态：此时 name 还未设置
response = await workflow.run(user_msg="What's my name?", ctx=ctx)
print(str(response))

# 通过工具设置 name，Agent 自动调用 set_name 函数
response2 = await workflow.run(user_msg="My name is Laurie", ctx=ctx)
print(str(response2))

# 直接从 ctx.store 读取状态值
state = await  ctx.store.get("state")
print("Name as stored in state: ",state["name"])

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.workflow import Context
# Tavily 搜索工具（需 API Key，比 DuckDuckGo 更稳定）
from llama_index.tools.tavily_research import TavilyToolSpec
import os
# 导入流式事件类型，用于区分不同的 Agent 行为
from llama_index.core.agent.workflow import (
    AgentInput,
    AgentOutput,
    ToolCall,
    ToolCallResult,
    AgentStream,
)


tavily_tool = TavilyToolSpec( api_key=envconfig.TAVILY_API_KEY)

workflow = AgentWorkflow.from_tools_or_functions(
    tavily_tool.to_tool_list(),
    llm=llm,
    system_prompt="You're a helpful assistant that can search the web for information."
)

async def main():
    handler = workflow.run(user_msg="深圳目前天气如何")

    # 流式输出：实时打印 Agent 的思考、工具调用、响应等事件
    async for event in handler.stream_events():
        print(event)
 
await main()

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.workflow import Context
from llama_index.tools.tavily_research import TavilyToolSpec
from llama_index.core.agent.workflow import (
    AgentInput,
    AgentOutput,
    ToolCall,
    ToolCallResult,
    AgentStream,
)


tavily_tool = TavilyToolSpec(api_key=envconfig.TAVILY_API_KEY)

workflow = AgentWorkflow.from_tools_or_functions(
    tavily_tool.to_tool_list(),
    llm=llm,
    system_prompt="You're a helpful assistant that can search the web for information."
)

async def main():
    handler = workflow.run(user_msg="深圳目前天气如何")

    # handle streaming output
    async for event in handler.stream_events():
        if isinstance(event, AgentStream):
            print(event.delta, end="", flush=True)   # 流式打印输出
        elif isinstance(event, AgentInput):
            # pass
            print("Agent input: ", event.input)  # the current input messages
            print("Agent name:", event.current_agent_name)  # the current agent name
        elif isinstance(event, AgentOutput):
            # pass
            print("Agent output: ", event.response)  # the current full response
            print("Tool calls made: ", event.tool_calls)  # the selected tool calls, if any
            print("Raw LLM response: ", event.raw)  # the raw llm api response
        elif isinstance(event, ToolCallResult):
            # pass
            print("Tool called: ", event.tool_name)  # the tool name
            print("Arguments to the tool: ", event.tool_kwargs)  # the tool kwargs
            print("Tool output: ", event.tool_output)  # the tool output            

    # print final output
    print("===========================")
    print(str(await handler))  # 最终打印一次完整输出
 
await main()

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.workflow import Context
# InputRequiredEvent: 这是一个“出站事件”，工具用来向外界大喊“我需要人类输入！”
# HumanResponseEvent: 这是一个“入站事件”，外界用来向工具回传“这就是人类的输入！”
from llama_index.core.workflow import (
    InputRequiredEvent,
    HumanResponseEvent,
)



# 定义一个需要人类确认的高危工具函数
# 必须传入 ctx: Context 对象，因为我们需要用它在异步环境里收发事件
async def dangerous_task(ctx: Context) -> str:
    """这是一个高危任务，在执行前必须获得人类的安全确认。"""

    # 1. 🚀 发送信号：往全局事件流里写一个 InputRequiredEvent。
    # 这会立刻被外部正在 stream_events() 的主线程捕获到。
    # ctx.write_event_to_stream(
    #     InputRequiredEvent(
    #         prefix="Are you sure you want to proceed? ", # 提示词
    #         user_name="Laurie",                          # 指定必须由特定的人确认
    #     )
    # )

    # # 2. 💤 挂起等待：原地进入异步深度睡眠。
    # # 程序执行到这里会停住，直到外部发回一个 HumanResponseEvent，且里面的 user_name 必须是 "Laurie"
    # # 只有条件满足，才会唤醒此函数，并将事件对象赋值给 response
    # response = await ctx.wait_for_event(
    #     HumanResponseEvent, requirements={"user_name": "Laurie"}
    # )
    question = "Are you sure you want to proceed? "
    response = await ctx.wait_for_event(
        HumanResponseEvent,
        waiter_id=question,
        waiter_event=InputRequiredEvent(
            prefix=question,
            user_name="Laurie",
        ),
        requirements={"user_name": "Laurie"},
    )


    # 3. 🎉 唤醒继续：拿到人类的回复，开始做逻辑判断
    # response.response 拿到的就是人类在键盘上输入的字符串
    if response.response.strip().lower() == "yes":
        return "Dangerous task completed successfully."  # 如果输入的是 yes，任务执行成功
    else:
        return "Dangerous task aborted."                  # 否则，安全中止任务


# 组装工作流：将 dangerous_task 工具和大模型绑定在一起
workflow = AgentWorkflow.from_tools_or_functions(
    [dangerous_task],
    llm=llm,
    system_prompt="You are a helpful assistant that can perform dangerous tasks.",
)


async def main():
    # 1. 启动工作流：传入用户的初始指令
    handler = workflow.run(user_msg="I want to proceed with the dangerous task.")

    # 2. 监听事件管道：用异步迭代器实时监听后台源源不断扔出来的事件
    async for event in handler.stream_events():

        
        # 3. 捕捉高危确认请求：如果管道里飘出了一个 InputRequiredEvent
        if isinstance(event, InputRequiredEvent):
            
            # 4. 键盘交互：调用 Python 阻塞式的 input()，在终端等待用户输入
            # 控制台会打印：Are you sure you want to proceed? 
            response = input(event.prefix)
            
            # 5. 反手回弹：把用户的输入封装成 HumanResponseEvent，强行塞回工作流的上下文中
            # 这行代码一执行，后台正在 await wait_for_event(...) 的高危工具会瞬间被唤醒！
            handler.ctx.send_event(
                HumanResponseEvent(
                    response=response,          # 传递你在终端敲下的 yes/no
                    user_name=event.user_name,  # 传递核对的用户名 (Laurie)

                )
            )

    # 6. 终局等待：等待大模型整理好工具返回的结果，输出最终文本
    response = await handler
    print(str(response))


# if __name__ == "__main__":
#     # 在标准的 .py 文件中，使用 asyncio 启动异步主函数
#     # 如果是在 Jupyter Notebook 里，可以直接用 `await main()` 运行
#     import asyncio
#     asyncio.run(main())

await main()