import asyncio
from python_a2a import AgentNetwork, AIAgentRouter, A2AClient, Task, Message, MessageRole, TextContent
import json
import uuid

from time import sleep


async def main():
    # 1.创建 AgentNetwork
    # 实例化对象
    network = AgentNetwork(name='TravelAgentNetwork')
    # 添加 agent  name可以自己定义
    network.add(name='WeatherAgent', agent_or_url='http://127.0.0.1:5008')
    network.add(name='TicketAgent', agent_or_url='http://127.0.0.1:5009')
    # 打印agent server信息
    for agent_info in network.list_agents():
        print(json.dumps(agent_info, indent=4, ensure_ascii=False))

    # 2.创建AIAgentRouter
    # 创建路由器
    router = AIAgentRouter(llm_client=A2AClient("http://127.0.0.1:5555"),
                           agent_network=network)

    # 3. 测试查询
    queries = [
        "帮我查下北京的天气",  # 应该路由到 WeatherAgent
        "预订一张从北京到上海的火车票",  # 应该路由到 TicketAgent
        "帮我查下北京的天气，并预订一张从北京到上海的火车票"  # 复合意图
    ]
    for query in queries:
        print(f"[主控日志] 用户查询: '{query}'")
        # 使用路由器选择agent
        agent_name, confidence = router.route_query(query)
        print(f"[主控日志] 匹配的agent: '{agent_name}', 匹配度: {confidence}")
        if agent_name:
            # 基于 agent_name 获取对应的 agent server
            agent_client = network.get_agent(agent_name)
            if agent_client:
                # Message用来存储具体的任务内容
                message = Message(content=TextContent(text=query), role=MessageRole.USER)
                # Task中存储和封装Message
                agent_task = Task(message=message.to_dict(), id="task-" + str(uuid.uuid4()))
                # 使用 send_task_async 实现异步调用
                try:
                    agent_result = await agent_client.send_task_async(agent_task)
                    print("[主控日志] agent返回结果:")
                    print(json.dumps(agent_result.to_dict(), indent=4, ensure_ascii=False))
                except Exception as e:
                    print(f"[主控日志] agent调用失败: '{e}'")




if __name__ == '__main__':
    asyncio.run(main())