微信群接入ComfyUI
# -*- coding: utf-8 -*-
import logging
import re
import time
import os
from queue import Empty
from threading import Thread
from wcferry import Wcf, WxMsg
from configuration import Config
from func_chatgpt import ChatGPT
from func_chengyu import cy
from func_news import News
from func_tigerbot import TigerBot
from job_mgmt import Job
import json
import requests
import base64
from translate import Translator
from comfyui_api import ComfyUI
class Robot(Job):
"""个性化自己的机器人
"""
def __init__(self, config: Config, wcf: Wcf, comfy_ui: ComfyUI) -> None:
self.wcf = wcf
self.config = config
self.LOG = logging.getLogger("Robot")
self.wxid = self.wcf.get_self_wxid()
self.allContacts = self.getAllContacts()
self.comfy_ui = comfy_ui
if self.config.TIGERBOT:
self.chat = TigerBot(self.config.TIGERBOT)
elif self.config.CHATGPT:
cgpt = self.config.CHATGPT
self.chat = ChatGPT(cgpt.get("key"), cgpt.get("api"), cgpt.get("proxy"), cgpt.get("prompt"))
else:
self.chat = None
# sd文生图api调用
def txt2Img(self, msg: WxMsg) -> bool:
if not msg.content.startswith('#img'):
return False
# 将中文提示词翻译成英文提示词
translator = Translator(to_lang="en", from_lang='zh')
en_text = translator.translate(msg.content.split('#img')[1])
self.LOG.info(f"prompt:{en_text}")
# 判断是否涉及nsfw
is_nsfw = self.nsfwFilter(en_text, msg.sender)
if is_nsfw:
self.sendTextMsg("命中NSFW：" + is_nsfw + "\n" + "已被移出群聊" + "\n" + "【请勿发送 NSFW 内容，请严格遵守群纪律以及相关法律法规，共同营造良好的群聊环境】", msg.roomid, msg.sender)
self.wcf.del_chatroom_members(msg.roomid, msg.sender)
return False
# 调用sd-webui文生图
# img = self.sd_webui(en_text)
# 调用ComfyUI文生图
img = self.comfy_ui.text_2_img(en_text)
# 将图片保存到当前目录下
with open("", 'wb') as fp:
fp.write(img)
# @ 微信群中那个要生成图片的男人，告诉他图片生成好了
self.sendTextMsg("图片生成成功，请查收！", msg.roomid, msg.sender)
# 将图片发送到微信群
img_path = os.path.dirname(os.path.abspath(__file__))
self.sendImg(img_path + "\\", msg.roomid)
return True
def sd_webui(self, en_text):
# 请求头
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/99.0.4844.84 Safari/537.36',
'Content-Type': 'application/json',
'Connection': 'keep-alive',
'Accept': '*/*',
'Accept-Encoding': 'gzip, deflate, br'
}
# 读取 中sd的文生图api的配置
sd = Config().SD
# 请求体
data = {
# 'prompt': ('prompt'),
'prompt': en_text,
'negative_prompt': sd.get('negative_prompt'),
'width': sd.get('width'),
'height': sd.get('height'),
'steps': sd.get('steps'),
'cfg_scale': sd.get('cfg_scale'),
'sampler_name': sd.get('sampler_name'),
'batch_size': sd.get('batch_size'),
'seed': sd.get('seed'),
'save_images': sd.get('save_images')
}
# 编码请求体
json_data = json.dumps(data).encode("utf-8")
print(json_data)
# 请求文生图api
response = requests.post(url=sd.get('url'), headers=headers, data=json_data)
# 获取文生图api的请求结果中的第一张图片
base64_img = json.loads(response.text).get("images")[0]
# 对获取到的图片进行base64解码
img = base64.b64decode(base64_img)
# with open("", 'w') as fp:
# (base64_img)
return img
# chatGLM
def chatGLM(self, msg: WxMsg) -> bool:
text = msg.content.split('#glm')[1]
# 请求头
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/99.0.4844.84 Safari/537.36',
'Content-Type': 'application/json',
'Connection': 'keep-alive',
'Accept': '*/*',
'Accept-Encoding': 'gzip, deflate, br'
}
# 读取 中 ChatGLM api的配置
chatglm = self.config.CHATGLM
# 请求体
data = {
'prompt': text,
'history': []
}
# 编码请求体
json_data = json.dumps(data).encode("utf-8")
print(json_data)
# 请求ChatGLM的api
response = requests.post(url=chatglm.get('url'), headers=headers, data=json_data)
# 将请求结果转为json并获取回答内容
answer = json.loads(response.text).get("response")
# @ 微信群中那个提问的男人，将ChatGLM的内容发给他
self.sendTextMsg(answer, msg.roomid, msg.sender)
return True
# 过滤NSFW内容
def nsfwFilter(self, en_text: str, sender: str) -> str:
nsfw = self.config.NSFW
nsfw_words = nsfw.get("word")
whitelilst = nsfw.get("whitelist")
en_arr = en_text.split(",")
if sender not in whitelilst:
for w in nsfw_words:
if w in en_arr:
self.LOG.info(f"【注意命中NSFW】：{w}：{sender}")
return w
return ""
def toAt(self, msg: WxMsg) -> bool:
"""处理被 @ 消息
:param msg: 微信消息结构
:return: 处理状态，`True` 成功，`False` 失败
"""
return self.toChitchat(msg)
def toChengyu(self, msg: WxMsg) -> bool:
"""
处理成语查询/接龙消息
:param msg: 微信消息结构
:return: 处理状态，`True` 成功，`False` 失败
"""
status = False
texts = re.findall(r"^([#|?|？])(.*)$", msg.content)
# [('#', '天天向上')]
if texts:
flag = texts[0][0]
text = texts[0][1]
if flag == "#": # 接龙
if cy.isChengyu(text):
rsp = cy.getNext(text)
if rsp:
self.sendTextMsg(rsp, msg.roomid)
status = True
elif flag in ["?", "？"]: # 查词
if cy.isChengyu(text):
rsp = cy.getMeaning(text)
if rsp:
self.sendTextMsg(rsp, msg.roomid)
status = True
return status
def toChitchat(self, msg: WxMsg) -> bool:
"""闲聊，接入 ChatGPT
"""
if not self.chat: # 没接 ChatGPT，固定回复
# rsp = "你@我干嘛？"
return False
else: # 接了 ChatGPT，智能回复
q = re.sub(r"@.*?[\u2005|\s]", "", msg.content).replace(" ", "")
rsp = self.chat.get_answer(q, (msg.roomid if msg.from_group() else msg.sender))
if rsp:
if msg.from_group():
self.sendTextMsg(rsp, msg.roomid, msg.sender)
else:
self.sendTextMsg(rsp, msg.sender)
return True
else:
self.LOG.error(f"无法从 ChatGPT 获得答案")
return False
def processMsg(self, msg: WxMsg) -> None:
"""当接收到消息的时候，会调用本方法。如果不实现本方法，则打印原始消息。
此处可进行自定义发送的内容,如通过 关键字自动获取当前天气信息，并发送到对应的群组@发送者
群号： 微信ID： 消息内容：
content = "xx天气信息为："
receivers =
(content, receivers, )
"""
# 群聊消息
if msg.from_group():
# 如果在群里被 @
if msg.roomid not in self.config.GROUPS: # 不在配置的响应的群列表里，忽略
return
if msg.is_at(self.wxid): # 被@
self.toAt(msg)
else: # 其他消息
if "外卖" in msg.content:
wm = Config().WM
mt = wm.get("mt")
ele = wm.get("ele")
content = "美团外卖红包 " + mt + "\n" + "饿了么外卖红包 " + ele
self.sendTextMsg(content, msg.roomid, msg.sender)
# 判断微信群消息是否包含指定关键字，从而触发文生图api的调用
# if "#国风" in or "#美女" in or "#国风美女" in :
# self.txt2Img(msg)
if "#img" in msg.content:
self.txt2Img(msg)
if "#glm" in msg.content:
self.chatGLM(msg)
else:
self.toChengyu(msg)
return # 处理完群聊信息，后面就不需要处理了
# 非群聊信息，按消息类型进行处理
if msg.type == 37: # 好友请求
self.autoAcceptFriendRequest(msg)
elif msg.type == 10000: # 系统信息
self.sayHiToNewFriend(msg)
elif msg.type == 0x01: # 文本消息
# 让配置加载更灵活，自己可以更新配置。也可以利用定时任务更新。
if msg.from_self():
if msg.content == "^更新$":
self.config.reload()
self.LOG.info("已更新")
else:
self.toChitchat(msg) # 闲聊
def onMsg(self, msg: WxMsg) -> int:
try:
# 只打印已配置的群聊中的信息
for wxid in self.config.GROUPS:
if wxid == msg.roomid:
self.LOG.info("sender:", msg.sender)
self.LOG.info("roomid:", msg.roomid)
self.LOG.info("type:", msg.type)
self.LOG.info(msg) # 打印信息
self.processMsg(msg)
except Exception as e:
self.LOG.error(e)
return 0
def enableRecvMsg(self) -> None:
self.wcf.enable_recv_msg(self.onMsg)
def enableReceivingMsg(self) -> None:
def innerProcessMsg(wcf: Wcf):
while wcf.is_receiving_msg():
try:
# 只打印已配置的群聊中的信息
msg = wcf.get_msg()
for wxid in self.config.GROUPS:
if wxid == msg.roomid:
self.LOG.info(f"sender:{msg.sender}")
self.LOG.info(f"roomid:{msg.roomid}")
self.LOG.info(f"type:{msg.type}")
self.LOG.info(msg)
self.processMsg(msg)
except Empty:
continue # Empty message
except Exception as e:
self.LOG.error(f"Receiving message error: {e}")
self.wcf.enable_receiving_msg()
Thread(target=innerProcessMsg, name="GetMessage", args=(self.wcf,), daemon=True).start()
def sendTextMsg(self, msg: str, receiver: str, at_list: str = "") -> None:
""" 发送消息
:param msg: 消息字符串
:param receiver: 接收人wxid或者群id
:param at_list: 要@的wxid, @所有人的wxid为：nofity@all
"""
# msg 中需要有 @ 名单中一样数量的 @
ats = ""
if at_list:
wxids = at_list.split(",")
for wxid in wxids:
# 这里偷个懒，直接 @昵称。有必要的话可以通过 里的 ChatRoom 表，解析群昵称
ats += f" @{self.allContacts.get(wxid, '')}"
# {msg}{ats} 表示要发送的消息内容后面紧跟@，例如 北京天气情况为：xxx @张三，微信规定需这样写，否则@不生效
if ats == "":
self.LOG.info(f"To {receiver}: {msg}")
self.wcf.send_text(f"{msg}", receiver, at_list)
else:
self.LOG.info(f"To {receiver}: {ats}\r{msg}")
self.wcf.send_text(f"{ats}\n\n{msg}", receiver, at_list)
def sendImg(self, path: str, receiver: str) -> None:
""" 发送消息
:param path: 图片路径
:param receiver: 接收人wxid或者群id
"""
self.LOG.info(f"To {receiver}: {path}")
tem = self.wcf.send_image(f"{path}", receiver)
def getAllContacts(self) -> dict:
"""
获取联系人（包括好友、公众号、服务号、群成员……）
格式: {"wxid": "NickName"}
"""
contacts = self.wcf.query_sql("", "SELECT UserName, NickName FROM Contact;")
return {contact["UserName"]: contact["NickName"] for contact in contacts}
def keepRunningAndBlockProcess(self) -> None:
"""
保持机器人运行，不让进程退出
"""
while True:
self.runPendingJobs()
time.sleep(1)
def autoAcceptFriendRequest(self, msg: WxMsg) -> None:
# 不自动通过好友
return
# try:
# xml = ()
# v3 = ["encryptusername"]
# v4 = ["ticket"]
# scene = int(["scene"])
# .accept_new_friend(v3, v4, scene)
#
# except Exception as e:
# (f"同意好友出错：{e}")
def sayHiToNewFriend(self, msg: WxMsg) -> None:
nickName = re.findall(r"你已添加了(.*)，现在可以开始聊天了。", msg.content)
if nickName:
# 添加了好友，更新好友列表
self.allContacts[msg.sender] = nickName[0]
self.sendTextMsg(f"Hi {nickName[0]}，我自动通过了你的好友请求。", msg.sender)
def newsReport(self) -> None:
receivers = self.config.NEWS
if not receivers:
return
news = News().get_important_news()
for r in receivers:
self.sendTextMsg(news, r)