wechat_service.py

from flask import Flask, request, make_response
import hashlib
import time
import xml.etree.ElementTree as ET
import os
import json
from openai import OpenAI
from dotenv import load_dotenv
from markdown import markdown
import re
import threading
import logging
from datetime import datetime
import asyncio
from concurrent.futures import ThreadPoolExecutor
import queue
import uuid
import base64
from Crypto.Cipher import AES
import struct
import random
import string
import requests

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('wechat_service.log'),
        logging.StreamHandler()
    ]
)

load_dotenv()

# Configuration
TOKEN = os.getenv('TOKEN')
ENCODING_AES_KEY = os.getenv('ENCODING_AES_KEY')
APPID = os.getenv('APPID')
APPSECRET = os.getenv('APPSECRET')
API_KEY = os.getenv("API_KEY")
BASE_URL = os.getenv("OPENAI_BASE_URL")
IMAGE_MODEL_URL = os.getenv("IMAGE_MODEL_URL")
IMAGE_MODEL_KEY = os.getenv("IMAGE_MODEL_KEY")

client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
executor = ThreadPoolExecutor(max_workers=10)

# Define tools for image generation
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "generate_image",
            "description": "Generate an image based on text description",
            "parameters": {
                "type": "object",
                "properties": {
                    "prompt": {
                        "type": "string",
                        "description": "The description of the image to generate"
                    }
                },
                "required": ["prompt"]
            }
        }
    }
]

class AccessTokenManager:
    def __init__(self):
        self._access_token = None
        self._expires_at = 0
        self._lock = threading.Lock()

    def get_token(self):
        with self._lock:
            now = time.time()
            # 提前5分钟刷新token，确保在调用时token都是有效的
            if self._access_token and now < (self._expires_at - 300):
                return self._access_token

            try:
                url = "https://api.weixin.qq.com/cgi-bin/token"
                params = {
                    "grant_type": "client_credential",
                    "appid": APPID,
                    "secret": APPSECRET
                }
                
                logging.info("开始获取新的access_token")
                response = requests.get(url, params=params)
                response.raise_for_status()
                result = response.json()
                
                if "access_token" not in result:
                    error_msg = f"获取access_token失败: {result}"
                    logging.error(error_msg)
                    raise ValueError(error_msg)
                
                self._access_token = result["access_token"]
                self._expires_at = now + result["expires_in"]
                logging.info("成功获取新的access_token")
                
                return self._access_token
                
            except Exception as e:
                error_msg = f"获取access_token时发生错误: {str(e)}"
                logging.error(error_msg)
                raise

    def refresh_token(self):
        with self._lock:
            self._access_token = None
            return self.get_token()
            
class WeChatCrypto:
    def __init__(self, key, app_id):
        self.key = base64.b64decode(key + '=')
        self.app_id = app_id

    def encrypt(self, text):
        random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
        text_bytes = text.encode('utf-8')
        msg_len = struct.pack('>I', len(text_bytes))
        message = random_str.encode('utf-8') + msg_len + text_bytes + self.app_id.encode('utf-8')
        pad_len = 32 - (len(message) % 32)
        message += chr(pad_len).encode('utf-8') * pad_len
        cipher = AES.new(self.key, AES.MODE_CBC, self.key[:16])
        encrypted = cipher.encrypt(message)
        return base64.b64encode(encrypted).decode('utf-8')

    def decrypt(self, encrypted_text):
        encrypted_data = base64.b64decode(encrypted_text)
        cipher = AES.new(self.key, AES.MODE_CBC, self.key[:16])
        decrypted = cipher.decrypt(encrypted_data)
        pad_len = decrypted[-1]
        if not isinstance(pad_len, int):
            pad_len = ord(pad_len)
        content = decrypted[16:-pad_len]
        msg_len = struct.unpack('>I', content[:4])[0]
        xml_content = content[4:msg_len + 4].decode('utf-8')
        app_id = content[msg_len + 4:].decode('utf-8')
        if app_id != self.app_id:
            raise ValueError('Invalid AppID')
        return xml_content

class AsyncResponse:
    def __init__(self):
        self.status = "processing"
        self.result = None
        self.error = None
        self.create_time = time.time()
        self.timeout = 3600
        self.response_type = "text"  # Can be "text" or "image"
        self.media_id = None  # For image responses

    def is_expired(self):
        return time.time() - self.create_time > self.timeout

class UserSession:
    def __init__(self):
        self.messages = [{"role": "system", "content": "你是HXIAO公众号的智能助手，这一个用来分享与学习人工智能的公众号，我们的目标是专注AI应用的简单研究与实践。致力于分享切实可行的技术方案，希望让复杂的技术变得简单易懂。也喜欢用通俗的语言来解释专业概念，让技术真正服务于每个学习者"}]
        self.pending_parts = []
        self.last_active = time.time()
        self.current_task = None
        self.response_queue = {}
        self.session_timeout = 3600

    def is_expired(self):
        return time.time() - self.last_active > self.session_timeout

    def cleanup_expired_tasks(self):
        expired_tasks = [
            task_id for task_id, response in self.response_queue.items()
            if response.is_expired()
        ]
        for task_id in expired_tasks:
            del self.response_queue[task_id]
            if self.current_task == task_id:
                self.current_task = None

class SessionManager:
    def __init__(self):
        self.sessions = {}
        self._lock = threading.Lock()
        self.crypto = WeChatCrypto(ENCODING_AES_KEY, APPID)

    def get_session(self, user_id):
        with self._lock:
            current_time = time.time()
            if user_id in self.sessions:
                session = self.sessions[user_id]
                if session.is_expired():
                    session = UserSession()
                else:
                    session.cleanup_expired_tasks()
            else:
                session = UserSession()
            session.last_active = current_time
            self.sessions[user_id] = session
            return session

    def clear_session(self, user_id):
        with self._lock:
            if user_id in self.sessions:
                self.sessions[user_id] = UserSession()

    def cleanup_expired_sessions(self):
        with self._lock:
            current_time = time.time()
            expired_users = [
                user_id for user_id, session in self.sessions.items()
                if session.is_expired()
            ]
            for user_id in expired_users:
                del self.sessions[user_id]
                logging.info(f"已清理过期会话: {user_id}")

def convert_markdown_to_wechat(md_text):
    if not md_text:
        return md_text
        
    md_text = re.sub(r'^# (.*?)$', r'【标题】\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^## (.*?)$', r'【子标题】\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^### (.*?)$', r'【小标题】\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'\*\*(.*?)\*\*', r'『\1』', md_text)
    md_text = re.sub(r'\*(.*?)\*', r'「\1」', md_text)
    md_text = re.sub(r'`(.*?)`', r'「\1」', md_text)
    md_text = re.sub(r'^\- ', '• ', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^\d\. ', '○ ', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'```[\w]*\n(.*?)```', r'【代码开始】\n\1\n【代码结束】', md_text, flags=re.DOTALL)
    md_text = re.sub(r'^> (.*?)$', r'▎\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^-{3,}$', r'—————————', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'\[(.*?)\]\((.*?)\)', r'\1（\2）', md_text)
    md_text = re.sub(r'\n{3,}', '\n\n', md_text)
    
    return md_text

def verify_signature(signature, timestamp, nonce, token):
    items = [token, timestamp, nonce]
    items.sort()
    temp_str = ''.join(items)
    hash_sha1 = hashlib.sha1(temp_str.encode('utf-8')).hexdigest()
    return hash_sha1 == signature

def verify_msg_signature(msg_signature, timestamp, nonce, token, encrypt_msg):
    """
    验证消息签名
    Args:
        msg_signature: 消息签名
        timestamp: 时间戳
        nonce: 随机数
        token: 验证令牌
        encrypt_msg: 加密的消息内容
    Returns:
        bool: 签名是否有效
    """
    items = [token, timestamp, nonce, encrypt_msg]
    items.sort()
    temp_str = ''.join(items)
    hash_sha1 = hashlib.sha1(temp_str.encode('utf-8')).hexdigest()
    return hash_sha1 == msg_signature


def parse_xml_message(xml_content):
    """
    解析微信XML消息，支持文本和图片消息类型
    """
    root = ET.fromstring(xml_content)
    message = {
        'from_user': root.find('FromUserName').text,
        'to_user': root.find('ToUserName').text,
        'create_time': root.find('CreateTime').text,
        'msg_type': root.find('MsgType').text,
        'msg_id': root.find('MsgId').text if root.find('MsgId') is not None else '',
        'msg_data_id': root.find('MsgDataId').text if root.find('MsgDataId') is not None else '',
        'idx': root.find('Idx').text if root.find('Idx') is not None else ''
    }
    
    if message['msg_type'] == 'text':
        message['content'] = root.find('Content').text if root.find('Content') is not None else ''
    elif message['msg_type'] == 'image':
        message['pic_url'] = root.find('PicUrl').text
        message['media_id'] = root.find('MediaId').text
    
    return message

def get_image_content(media_id):
    """
    通过微信接口获取图片内容
    """
    try:
        access_token = token_manager.get_token()
        url = f'https://api.weixin.qq.com/cgi-bin/media/get?access_token={access_token}&media_id={media_id}'
        
        logging.info(f"开始下载图片，media_id: {media_id}")
        response = requests.get(url)
        
        if response.headers.get('Content-Type') == 'text/plain':
            # 如果返回JSON错误信息
            error_info = response.json()
            if error_info.get('errcode') == 40001:
                # access_token过期，刷新后重试
                logging.info("access_token已过期，正在刷新并重试")
                access_token = token_manager.refresh_token()
                url = f'https://api.weixin.qq.com/cgi-bin/media/get?access_token={access_token}&media_id={media_id}'
                response = requests.get(url)
        
        response.raise_for_status()
        return response.content
        
    except Exception as e:
        logging.error(f"获取图片内容失败: {str(e)}")
        raise


def generate_response_xml(to_user, from_user, content, response_type='text', media_id=None, encrypt_type='aes'):
    timestamp = str(int(time.time()))
    nonce = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
    
    if response_type == 'image' and media_id:
        xml_content = f'''
        <xml>
            <ToUserName><![CDATA[{to_user}]]></ToUserName>
            <FromUserName><![CDATA[{from_user}]]></FromUserName>
            <CreateTime>{timestamp}</CreateTime>
            <MsgType><![CDATA[image]]></MsgType>
            <Image>
                <MediaId><![CDATA[{media_id}]]></MediaId>
            </Image>
        </xml>
        '''
    else:
        formatted_content = convert_markdown_to_wechat(content)
        xml_content = f'''
        <xml>
            <ToUserName><![CDATA[{to_user}]]></ToUserName>
            <FromUserName><![CDATA[{from_user}]]></FromUserName>
            <CreateTime>{timestamp}</CreateTime>
            <MsgType><![CDATA[text]]></MsgType>
            <Content><![CDATA[{formatted_content}]]></Content>
        </xml>
        '''
    
    if encrypt_type == 'aes':
        encrypted = session_manager.crypto.encrypt(xml_content)
        signature_list = [TOKEN, timestamp, nonce, encrypted]
        signature_list.sort()
        msg_signature = hashlib.sha1(''.join(signature_list).encode('utf-8')).hexdigest()
        
        response_xml = f'''
        <xml>
            <Encrypt><![CDATA[{encrypted}]]></Encrypt>
            <MsgSignature><![CDATA[{msg_signature}]]></MsgSignature>
            <TimeStamp>{timestamp}</TimeStamp>
            <Nonce><![CDATA[{nonce}]]></Nonce>
        </xml>
        '''
    else:
        response_xml = xml_content
    
    response = make_response(response_xml)
    response.content_type = 'application/xml'
    return response
    
# 创建全局的token管理器实例
token_manager = AccessTokenManager()

def upload_image_to_wechat(image_data):
    """上传图片到微信服务器并获取media_id"""
    try:
        access_token = token_manager.get_token()
        upload_url = f'https://api.weixin.qq.com/cgi-bin/media/upload?access_token={access_token}&type=image'
        files = {'media': ('image.jpg', image_data, 'image/jpeg')}
        
        logging.info("开始上传图片到微信服务器")
        response = requests.post(upload_url, files=files)
        response.raise_for_status()
        result = response.json()
        
        if 'media_id' not in result:
            if 'errcode' in result and result['errcode'] == 40001:
                # access_token 可能过期，尝试刷新并重试
                logging.info("access_token已过期，正在刷新并重试")
                access_token = token_manager.refresh_token()
                upload_url = f'https://api.weixin.qq.com/cgi-bin/media/upload?access_token={access_token}&type=image'
                response = requests.post(upload_url, files=files)
                response.raise_for_status()
                result = response.json()
            
            if 'media_id' not in result:
                error_msg = f"上传图片失败: {result}"
                logging.error(error_msg)
                raise ValueError(error_msg)
        
        logging.info(f"图片上传成功，获取到media_id")
        return result['media_id']
        
    except Exception as e:
        error_msg = f"上传图片过程中发生错误: {str(e)}"
        logging.error(error_msg)
        raise

def process_long_running_task(messages, message_type='text', image_data=None):
    """
    处理长时间运行的任务，支持文本对话和图片识别
    """
    try:
        logging.info(f"开始调用AI服务，消息类型: {message_type}")
        
        if message_type == 'image':
            # 图片识别逻辑保持不变
            try:
                image_content = get_image_content(image_data['media_id'])
                image_base64 = base64.b64encode(image_content).decode('utf-8')
                
                image_messages = [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "请详细描述这张图片中的内容，包括主要对象、场景、活动等关键信息"},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_base64}"
                                }
                            }
                        ]
                    }
                ]
                
                logging.info("开始调用图像识别模型")
                image_response = client.chat.completions.create(
                    model="gpt-4.1-mini",
                    messages=image_messages,
                    max_tokens=300,
                    timeout=60
                )
                logging.info("图像识别完成")
                
                if not image_response.choices:
                    raise Exception("图像识别服务未返回有效结果")
                    
                return {
                    "type": "text",
                    "content": image_response.choices[0].message.content
                }
                
            except Exception as e:
                logging.error(f"图像识别过程中发生错误: {str(e)}")
                raise
                
        else:
            # 处理文本消息
            try:
                logging.info("开始处理文本消息")
                response = client.chat.completions.create(
                    model="gpt-4.1-mini",
                    messages=messages,
                    tools=TOOLS,
                    tool_choice="auto",
                    timeout=60
                )
                
                # 检查是否需要生成图片
                if response.choices[0].message.tool_calls:
                    tool_call = response.choices[0].message.tool_calls[0]
                    if tool_call.function.name == "generate_image":
                        try:
                            logging.info("检测到图片生成请求")
                            args = json.loads(tool_call.function.arguments)
                            
                            # 使用新的DALL-E 3 API进行图片生成
                            image_generation_response = requests.post(
                                "https://api1.oaipro.com/v1/images/generations",
                                headers={
                                    'Content-Type': 'application/json',
                                    'Authorization': f'Bearer {API_KEY}'
                                },
                                json={
                                    "model": "dall-e-3",
                                    "prompt": args['prompt'],
                                    "n": 1,
                                    "size": "1024x1024"
                                },
                                timeout=60
                            )
                            image_generation_response.raise_for_status()
                            generation_result = image_generation_response.json()
                            
                            if 'data' not in generation_result or not generation_result['data']:
                                raise Exception("图片生成服务未返回有效结果")
                            
                            # 获取生成的图片URL
                            image_url = generation_result['data'][0]['url']
                            
                            # 下载生成的图片
                            img_response = requests.get(image_url, timeout=30)
                            img_response.raise_for_status()
                            
                            # 上传图片到微信服务器
                            media_id = upload_image_to_wechat(img_response.content)
                            
                            return {
                                "type": "image",
                                "media_id": media_id
                            }
                            
                        except requests.exceptions.RequestException as e:
                            logging.error(f"图片生成过程中发生网络错误: {str(e)}")
                            raise
                        except Exception as e:
                            logging.error(f"图片生成过程中发生错误: {str(e)}")
                            raise
                
                # 返回文本响应
                return {
                    "type": "text",
                    "content": response.choices[0].message.content
                }
                
            except requests.exceptions.RequestException as e:
                logging.error(f"处理文本消息时发生网络错误: {str(e)}")
                raise
            except Exception as e:
                logging.error(f"处理文本消息时发生错误: {str(e)}")
                raise
                
    except Exception as e:
        logging.error(f"API调用错误: {str(e)}")
        raise

def handle_async_task(session, task_id, messages=None, message_type='text', message_data=None):
    """
    处理异步任务，支持文本对话和图片识别
    """
    try:
        logging.info(f"开始处理异步任务: {task_id}, 类型: {message_type}")
        
        if task_id not in session.response_queue:
            return
            
        if message_type == 'image':
            result = process_long_running_task(None, 'image', message_data)
        else:
            result = process_long_running_task(messages)
            
        if task_id in session.response_queue and not session.response_queue[task_id].is_expired():
            session.response_queue[task_id].status = "completed"
            session.response_queue[task_id].response_type = result.get("type", "text")
            
            if result["type"] == "image":
                session.response_queue[task_id].media_id = result["media_id"]
                session.response_queue[task_id].result = None
            else:
                session.response_queue[task_id].result = result["content"]
                
            if messages and result["type"] == "text":
                messages.append({"role": "assistant", "content": result["content"]})
                
    except Exception as e:
        logging.error(f"异步任务处理失败: {str(e)}")
        if task_id in session.response_queue:
            session.response_queue[task_id].status = "failed"
            session.response_queue[task_id].error = str(e)


def generate_initial_response():
    return "您的请求正在处理中，请回复'查询'获取结果（生图需要时间）"

def split_message(message, max_length=500):
    """
    将长消息分割成多个部分
    Args:
        message: 需要分割的消息
        max_length: 每部分的最大长度
    Returns:
        list: 分割后的消息部分列表
    """
    return [message[i:i+max_length] for i in range(0, len(message), max_length)]

def append_status_message(content, has_pending_parts=False, is_processing=False):
    """
    添加状态消息到响应内容
    Args:
        content: 原始内容
        has_pending_parts: 是否有待发送的部分
        is_processing: 是否正在处理中
    Returns:
        str: 添加了状态信息的内容
    """
    if "您的请求正在处理中" in content:
        return content + "\n\n-------------------\n发送'新对话'开始新的对话"
        
    status_message = "\n\n-------------------"
    if is_processing:
        status_message += "\n请回复'查询'获取结果"
    elif has_pending_parts:
        status_message += "\n当前消息已截断，发送'继续'查看后续内容"
    status_message += "\n发送'新对话'开始新的对话"
    return content + status_message

session_manager = SessionManager()

def cleanup_sessions():
    while True:
        time.sleep(3600)  # 每小时清理一次
        try:
            session_manager.cleanup_expired_sessions()
        except Exception as e:
            logging.error(f"清理会话时出错: {str(e)}")