from python_a2a import AIAgentRouter, AgentNetwork
from langchain_openai import ChatOpenAI
from agent_learn.config import Config

conf=Config()

# 创建网络
network = AgentNetwork(name="MyNetwork")
network.add("TicketAgent", "http://127.0.0.1:5010")

# 创建模型
llm = ChatOpenAI(base_url=conf.base_url,
                 api_key=conf.api_key,
                 model=conf.model_name,
                 temperature=0.1)

# 创建路由器
router = AIAgentRouter(llm_client=llm, agent_network=network)
# 输入查询，让 路由器 找到最合适的代理
question = "定一张北京去上海的火车票"
agent_name, confidence = router.route_query(question)
print(f"{agent_name}")

# 调用 通过 name 获取对应的 agent，然后使用 ask 方法进行提问
client = network.get_agent(name=agent_name)
print(client.ask(question))