from python_a2a import AIAgentRouter, AgentNetwork, Message, TextContent, MessageRole
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)

# 用一个小适配类，把 A2A 的 send_message 转成 LangChain 的 invoke。
class LangChainRouterClient:
    def send_message(self, message):
        answer = llm.invoke(message.content.text)
        return Message(content=TextContent(text=answer.content), role=MessageRole.AGENT)

# 创建路由器
router = AIAgentRouter(llm_client=LangChainRouterClient(), agent_network=network)
agent_name, confidence = router.route_query("预订票")
print(agent_name, confidence)