161 lines
5.7 KiB
Python
161 lines
5.7 KiB
Python
import os
|
||
import time
|
||
import logging
|
||
from aiohttp import web
|
||
from wechatpy.crypto import WeChatCrypto
|
||
from wechatpy.exceptions import InvalidSignatureException
|
||
from wechatpy import parse_message
|
||
from wechatpy.replies import TextReply
|
||
from anthropic import AsyncAnthropic
|
||
|
||
# --- 配置区域 ---
|
||
CONFIG = {
|
||
"WOA_TOKEN": os.environ['WOA_TOKEN'], # 微信公众号后台设置的 Token
|
||
"WOA_APPID": os.environ['WOA_APPID'], # kaiyunayun 公众号 AppID
|
||
"WOA_AESKEY": os.environ['WOA_AESKEY'] # 如果开启加密模式则填写，否则留空
|
||
}
|
||
|
||
# 初始化日志
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 初始化客户端
|
||
claude_client = AsyncAnthropic(api_key=CONFIG["ANTHROPIC_API_KEY"])
|
||
crypto = None
|
||
if CONFIG["WECHAT_ENCODING_AES_KEY"]:
|
||
crypto = WeChatCrypto(CONFIG["WECHAT_TOKEN"], CONFIG["WECHAT_ENCODING_AES_KEY"], CONFIG["WECHAT_APP_ID"])
|
||
|
||
# --- 内存数据库 (生产环境请替换为 Redis) ---
|
||
# 结构: { openid: [ {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."} ] }
|
||
user_sessions = {}
|
||
|
||
def get_system_prompt():
|
||
return "你是一个运行在微信公众号里的智能助手，回答要简洁、友好，适合手机端阅读。"
|
||
|
||
async def get_claude_reply(openid, user_content):
|
||
"""调用 Claude API 获取回复"""
|
||
|
||
# 1. 初始化或获取用户会话历史
|
||
if openid not in user_sessions:
|
||
user_sessions[openid] = []
|
||
|
||
history = user_sessions[openid]
|
||
|
||
# 2. 添加当前用户消息
|
||
history.append({"role": "user", "content": user_content})
|
||
|
||
# 3. 限制历史记录长度 (防止 token 超限)
|
||
if len(history) > CONFIG["MAX_HISTORY"] * 2:
|
||
history = history[-CONFIG["MAX_HISTORY"] * 2:]
|
||
user_sessions[openid] = history
|
||
|
||
try:
|
||
# 4. 调用 Claude
|
||
response = await claude_client.messages.create(
|
||
model="claude-3-5-sonnet-20241022",
|
||
max_tokens=1024,
|
||
system=get_system_prompt(),
|
||
messages=history
|
||
)
|
||
|
||
ai_reply = response.content[0].text
|
||
|
||
# 5. 保存 AI 回复到历史
|
||
history.append({"role": "assistant", "content": ai_reply})
|
||
|
||
return ai_reply
|
||
|
||
except Exception as e:
|
||
logger.error(f"Claude API Error: {e}")
|
||
return "抱歉，大脑暂时短路了，请稍后再试。"
|
||
|
||
# --- HTTP 路由处理 ---
|
||
|
||
async def handle_wechat(request):
|
||
"""处理微信回调 (GET 验证 & POST 消息)"""
|
||
query = request.query
|
||
|
||
# 1. GET 请求：微信服务器验证
|
||
if request.method == "GET":
|
||
signature = query.get("signature", "")
|
||
timestamp = query.get("timestamp", "")
|
||
nonce = query.get("nonce", "")
|
||
echostr = query.get("echostr", "")
|
||
|
||
try:
|
||
# 如果开启了加密模式，使用 crypto.check_signature，否则手动验证
|
||
if crypto:
|
||
crypto.check_signature(signature, timestamp, nonce)
|
||
else:
|
||
# 手动验证逻辑 (简单版)
|
||
import hashlib
|
||
tmp_list = sorted([CONFIG["WECHAT_TOKEN"], timestamp, nonce])
|
||
tmp_str = "".join(tmp_list)
|
||
hash_str = hashlib.sha1(tmp_str.encode('utf-8')).hexdigest()
|
||
if hash_str != signature:
|
||
raise InvalidSignatureException()
|
||
|
||
return web.Response(text=echostr)
|
||
except InvalidSignatureException:
|
||
logger.warning("Invalid signature from WeChat")
|
||
return web.Response(text="Invalid signature", status=403)
|
||
|
||
# 2. POST 请求：接收用户消息
|
||
elif request.method == "POST":
|
||
body = await request.read()
|
||
|
||
# 如果开启加密，先解密
|
||
if crypto:
|
||
msg_signature = query.get("msg_signature", "")
|
||
timestamp = query.get("timestamp", "")
|
||
nonce = query.get("nonce", "")
|
||
try:
|
||
body = crypto.decrypt_message(body, msg_signature, timestamp, nonce)
|
||
except Exception as e:
|
||
logger.error(f"Decrypt error: {e}")
|
||
return web.Response(text="error")
|
||
|
||
# 解析消息
|
||
try:
|
||
msg = parse_message(body)
|
||
except Exception as e:
|
||
logger.error(f"Parse message error: {e}")
|
||
return web.Response(text="success") # 即使解析失败也返回 success 避免微信重试风暴
|
||
|
||
# 只处理文本消息
|
||
if msg.type == 'text':
|
||
openid = msg.source
|
||
user_content = msg.content
|
||
|
||
logger.info(f"User {openid}: {user_content}")
|
||
|
||
# 异步调用 Claude
|
||
reply_text = await get_claude_reply(openid, user_content)
|
||
|
||
# 构建回复
|
||
reply = TextReply(content=reply_text, message=msg)
|
||
response_xml = reply.render()
|
||
|
||
# 如果开启加密，需要加密返回
|
||
if crypto:
|
||
response_xml = crypto.encrypt_message(response_xml, nonce, timestamp)
|
||
|
||
return web.Response(text=response_xml, content_type='application/xml')
|
||
|
||
# 非文本消息或事件推送，直接返回 success 避免重试
|
||
return web.Response(text="success")
|
||
|
||
return web.Response(status=405)
|
||
|
||
# --- 启动服务 ---
|
||
|
||
async def init_app():
|
||
app = web.Application()
|
||
app.router.add_route('*', '/wechat', handle_wechat) # 路径需与公众号后台配置一致
|
||
return app
|
||
|
||
if __name__ == '__main__':
|
||
app = init_app()
|
||
logger.info(f"Starting WeChat Bot on port {CONFIG['PORT']}...")
|
||
web.run_app(app, host='0.0.0.0', port=CONFIG['PORT'])
|