examples/nxlink_bot.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
import logging

import dingtalk_stream
from dingtalk_stream import AckMessage

from project_settings import project_path
from project_settings import environment
from toolbox.agent_x.question_answer import AgentX


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--client_id",
        default=environment.get("client_id"),
        type=str,
    )
    parser.add_argument(
        "--client_secret",
        default=environment.get("client_secret"),
        type=str,
    )
    parser.add_argument(
        "--agent_x_api_key",
        default=environment.get("agent_x_api_key", default=None),
        type=str
    )
    parser.add_argument(
        "--agent_x_agent_name",
        default=environment.get("agent_x_agent_name", default=None),
        type=str
    )
    args = parser.parse_args()
    return args


class EchoTextHandler(dingtalk_stream.ChatbotHandler):
    def __init__(self):
        super(EchoTextHandler, self).__init__()

    async def process(self, callback: dingtalk_stream.CallbackMessage):
        incoming_message = dingtalk_stream.ChatbotMessage.from_dict(callback.data)
        text = incoming_message.text.content.strip()
        self.reply_text(text, incoming_message)
        return AckMessage.STATUS_OK, "OK"


class AgentXNXLinkHandler(dingtalk_stream.ChatbotHandler):
    def __init__(self, agent: AgentX):
        super(AgentXNXLinkHandler, self).__init__()
        self.agent = agent

    async def process(self, callback: dingtalk_stream.CallbackMessage):
        incoming_message = dingtalk_stream.ChatbotMessage.from_dict(callback.data)
        text = incoming_message.text.content.strip()

        response = self.agent.question_answer(
            question=text,
            streaming=True
        )
        answer = response["answer"]
        answer = answer.split("{\"message\":")[0]

        self.reply_text(answer, incoming_message)
        return AckMessage.STATUS_OK, "OK"


def main():
    args = get_args()

    agent = AgentX(
        api_key=args.agent_x_api_key,
        agent_name=args.agent_x_agent_name,
    )

    agent_x_nxlink_handler = AgentXNXLinkHandler(agent=agent)

    client = dingtalk_stream.DingTalkStreamClient(
        credential=dingtalk_stream.Credential(
            client_id=args.client_id,
            client_secret=args.client_secret,
        )
    )

    client.register_callback_handler(
        dingtalk_stream.chatbot.ChatbotMessage.TOPIC,
        agent_x_nxlink_handler
    )
    client.start_forever()
    return


if __name__ == '__main__':
    main()