wechat-server/web.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import time
from collections import defaultdict
from functools import lru_cache
from xml.etree.ElementTree import fromstring

import requests
import uvicorn
from fastapi import FastAPI, Request, Response
from WXBizMsgCrypt3 import WXBizMsgCrypt

from edu_assistant.learning_tasks import QaTask

instruction = """
Act as a c++ professional to answer student aged 5-10 questions. Answer properly and politely.
"""
task = QaTask(instruction=instruction)

app = FastAPI()

TEXT_RESPONSE_TEMPLATE = """
<xml>
   <ToUserName>{to_username}</ToUserName>
   <FromUserName>{from_username}</FromUserName>
   <CreateTime>{create_time}</CreateTime>
   <MsgType>text</MsgType>
   <Content>{content}</Content>
</xml>
"""

WECHAT_TOKEN = os.environ.get("WECHAT_TOKEN")
WECHAT_AESKEY = os.environ.get("WECHAT_AESKEY")
WECHAT_CORPID = os.environ.get("WECHAT_CORPID")
WECHAT_SECRET = os.environ.get("WECHAT_SECRET")
CODEDOG_PORT = int(os.environ.get("CODEDOG_PORT", 32167))
wxcpt = WXBizMsgCrypt(WECHAT_TOKEN, WECHAT_AESKEY, WECHAT_CORPID)

# 先用个简单的字典存一下，后面可以考虑用 redis
last_reply_time = dict()
server_start_time = int(time.time())


def access_token():
    url_base = "https://qyapi.weixin.qq.com/cgi-bin/gettoken?corpid=%s&corpsecret=%s"
    resp = requests.get(url_base % (WECHAT_CORPID, WECHAT_SECRET)).json()
    return resp.get('access_token', '')


def get_chat_history(access_token: str, message_token: str):
    url_base = "https://qyapi.weixin.qq.com/cgi-bin/kf/sync_msg?access_token=" + access_token
    body = {"token": message_token}
    resp = requests.post(url_base, json=body)
    """
    [{
      "msgid": "",
      "open_kfid": "",
      "external_userid": "",
      "send_time": 1691854816,
      "origin": 3,
      "msgtype": "text",
      "text": {
        "content": "哈哈哈哈哈啊"
      }
    }]
    """
    # print("history", json.dumps(resp.json(), indent=2, ensure_ascii=False))
    return resp.json()


def reply(external_user_id, open_kfid, content):
    url_base = "https://qyapi.weixin.qq.com/cgi-bin/kf/send_msg?debug=1&access_token=" + access_token()
    body = {
        "touser": external_user_id,
        "open_kfid": open_kfid,
        "msgtype": "text",
        "text": {
            "content": content
        }
    }
    resp = requests.post(url_base, json=body).json()
    return resp.get("errcode", 1) == 0


def extract_messages_blocks(history: list):
    text_message_blocks = []

    # 这里收集一下 我们能处理的消息
    for message_block in history:
        open_kfid = message_block.get("open_kfid", None)
        external_userid = message_block.get("external_userid", None)
        if open_kfid is None or external_userid is None:
            print("open_kfid or external_userid is None")
            continue
        print("message_block", message_block)
        if message_block.get("msgtype", None) == "text":
            content = message_block["text"]["content"]
            # 限制一下要大于服务启动的时间
            if message_block["send_time"] > last_reply_time.get(external_userid, server_start_time):
                text_message_blocks.append(message_block)
        if message_block.get("msgtype", None) == "image":
            content = message_block["image"]["media_id"]
            print("image block", content)

    grouped_blocks = defaultdict(list)

    if len(text_message_blocks) == 0:
        return grouped_blocks

    for message_block in text_message_blocks:
        grouped_blocks[message_block["external_userid"]].append(message_block)

    return grouped_blocks

# 企业微信同一条用户的消息会发三次回调，内容都一样
# 这里加个 lru_cache 挡一挡，纯属偷懒


@lru_cache()
def handle_message(message_token: str):

    history = get_chat_history(access_token=access_token(),
                               message_token=message_token)
    # 实际上这里能拿到所有的用户发来的消息，所以这里要做一下分组
    # TODO 翻页
    grouped_blocks = extract_messages_blocks(history["msg_list"])
    if len(grouped_blocks) == 0:
        return

    for external_userid, text_message_blocks in grouped_blocks.items():
        content = "\n".join([message_block["text"]["content"] for message_block in text_message_blocks])
        open_kfid = text_message_blocks[-1].get("open_kfid", None)
        last_reply_time[text_message_blocks[-1]["external_userid"]] = text_message_blocks[-1]["send_time"]

        # print(open_kfid, external_userid, content)
        result = task.ask(content, session=False)
        print("handle_message", result)
        """
        {
            "input": "",
            "chat_history": "",
            "text": ""
        }
        """
        reply(external_userid, open_kfid, result["text"])


@app.get("/")
async def verify(msg_signature: str,
                 timestamp: str,
                 nonce: str,
                 echostr: str):
    '''
    验证配置是否成功，处理get请求
    :param msg_signature:
    :param timestamp:
    :param nonce:
    :param echostr:
    :return:
    '''
    ret, sEchoStr = wxcpt.VerifyURL(msg_signature, timestamp, nonce, echostr)
    if ret == 0:
        return Response(content=sEchoStr.decode('utf-8'))
    else:
        print(sEchoStr)


@app.post("/")
async def recv(msg_signature: str,
               timestamp: str,
               nonce: str,
               request: Request):
    '''
    接收用户消息，可进行被动响应
    :param msg_signature:
    :param timestamp:
    :param nonce:
    :param request:
    :return:
    '''
    body = await request.body()
    ret, msg = wxcpt.DecryptMsg(body.decode(
        'utf-8'), msg_signature, timestamp, nonce)
    decrypt_data = {}
    for node in list(fromstring(msg.decode('utf-8'))):
        decrypt_data[node.tag] = node.text
    # 解析后得到的decrypt_data:
    # {"ToUserName":"企业号", "FromUserName":"发送者用户名", "CreateTime":"发送时间",
    # "Content":"用户发送的内容", "MsgId":"唯一id，需要针对此id做出响应", "AagentID": "应用id"}
    print("decrypt_data", decrypt_data)

    message_token = decrypt_data["Token"]
    # 真正进行回复的地方
    handle_message(message_token)

    # 这里只需要响应一下回调函数即可，不需要返回数据
    return Response(content="success")

    # resp_data = TEXT_RESPONSE_TEMPLATE.format(to_username=decrypt_data.get("ToUserName", ""),
    #                                           from_username=decrypt_data.get(
    #                                               "FromUserName", ""),
    #                                           create_time=decrypt_data.get(
    #                                               "CreateTime", ""),
    #                                           content="帅得一逼",)
    # ret, send_msg = wxcpt.EncryptMsg(sReplyMsg=resp_data, sNonce=nonce)
    # if ret == 0:
    #     return Response(content=send_msg)
    # else:
    #     print(send_msg)


if __name__ == "__main__":
    uvicorn.run("web:app", port=CODEDOG_PORT, host='0.0.0.0', reload=False)