Update app.py

app.py

@@ -1,3 +1,48 @@
+Hugging Face's logo
+Hugging Face
+Models
+Datasets
+Spaces
+Posts
+Docs
+Enterprise
+Pricing
+
+
+
+Spaces:
+
+mistpe
+/
+wewe
+
+
+like
+1
+
+App
+Files
+Community
+Settings
+wewe
+/
+app.py
+
+mistpe's picture
+mistpe
+Update app.py
+c2c436e
+verified
+5 days ago
+raw
+
+Copy download link
+history
+blame
+edit
+delete
+
+25 kB
 from flask import Flask, request, make_response
 import hashlib
 import time
@@ -247,22 +292,75 @@ def verify_signature(signature, timestamp, nonce, token):
     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)
-    return {
-        'content': root.find('Content').text if root.find('Content') is not None else '',
+    message = {
         'from_user': root.find('FromUserName').text,
         'to_user': root.find('ToUserName').text,
-        'msg_id': root.find('MsgId').text if root.find('MsgId') is not None else '',
         'create_time': root.find('CreateTime').text,
-        'msg_type': root.find('MsgType').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()))
@@ -350,96 +448,213 @@ def upload_image_to_wechat(image_data):
         logging.error(error_msg)
         raise
 
-def process_long_running_task(messages):
+def process_long_running_task(messages, message_type='text', image_data=None):
+    """
+    处理长时间运行的任务，支持文本对话和图片识别
+    Args:
+        messages: 消息历史记录列表
+        message_type: 消息类型，'text'或'image'
+        image_data: 图片相关数据，包含media_id等信息
+    Returns:
+        dict: 包含处理结果的字典
+    """
     try:
-        logging.info("开始调用AI服务")
-        response = client.chat.completions.create(
-            model="o3-mini",
-            messages=messages,
-            tools=TOOLS,
-            tool_choice="auto",
-            timeout=60
-        )
-        logging.info("AI服务响应成功")
-
-        if response.choices[0].message.tool_calls:
-            logging.info("检测到tool调用")
-            tool_call = response.choices[0].message.tool_calls[0]
-            if tool_call.function.name == "generate_image":
-                logging.info("开始处理图片生成请求")
-                args = json.loads(tool_call.function.arguments)
-                image_response = requests.post(
-                    IMAGE_MODEL_URL,
-                    headers={
-                        'Content-Type': 'application/json',
-                        'Authorization': f'Bearer {IMAGE_MODEL_KEY}'
-                    },
-                    json={
-                        "model": "grok-latest-image",
-                        "messages": [{
-                            "role": "user",
-                            "content": args['prompt']
-                        }]
+        logging.info(f"开始调用AI服务，消息类型: {message_type}")
+        
+        if message_type == 'image':
+            try:
+                # 获取图片内容
+                image_content = get_image_content(image_data['media_id'])
+                # 将图片内容转换为base64
+                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-4o-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 requests.exceptions.RequestException as e:
+                logging.error(f"获取或处理图片时发生网络错误: {str(e)}")
+                raise
+            except Exception as e:
+                logging.error(f"图像识别过程中发生错误: {str(e)}")
+                raise
+                
+        else:
+            # 处理文本消息
+            try:
+                logging.info("开始处理文本消息")
+                response = client.chat.completions.create(
+                    model="o3-mini",
+                    messages=messages,
+                    tools=TOOLS,
+                    tool_choice="auto",
+                    timeout=60
                 )
-                image_response.raise_for_status()
-                result = image_response.json()
-                logging.info("图片生成成功，准备下载图片")
-
-                # 从markdown格式中提取URL
-                markdown_content = result['choices'][0]['message']['content']
-                image_url = re.search(r'\!\[image\]\((.*?)\)', markdown_content).group(1)
-                logging.info(f"提取到图片URL: {image_url}")
                 
-                img_response = requests.get(image_url)
-                img_response.raise_for_status()
-                media_id = upload_image_to_wechat(img_response.content)
+                # 检查是否需要生成图片
+                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)
+                            
+                            # 调用图片生成服务
+                            image_generation_response = requests.post(
+                                IMAGE_MODEL_URL,
+                                headers={
+                                    'Content-Type': 'application/json',
+                                    'Authorization': f'Bearer {IMAGE_MODEL_KEY}'
+                                },
+                                json={
+                                    "model": "grok-latest-image",
+                                    "messages": [{
+                                        "role": "user",
+                                        "content": args['prompt']
+                                    }]
+                                },
+                                timeout=60
+                            )
+                            image_generation_response.raise_for_status()
+                            generation_result = image_generation_response.json()
+                            
+                            if 'choices' not in generation_result or not generation_result['choices']:
+                                raise Exception("图片生成服务未返回有效结果")
+                                
+                            # 从markdown格式中提取URL
+                            markdown_content = generation_result['choices'][0]['message']['content']
+                            image_url_match = re.search(r'\!\[image\]\((.*?)\)', markdown_content)
+                            if not image_url_match:
+                                raise Exception("无法从响应中提取图片URL")
+                                
+                            image_url = image_url_match.group(1)
+                            
+                            # 下载生成的图片
+                            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": "image",
-                    "media_id": media_id
+                    "type": "text",
+                    "content": response.choices[0].message.content
                 }
-
-        logging.info("返回文本响应")
-        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):
+
+def handle_async_task(session, task_id, messages=None, message_type='text', message_data=None):
+    """
+    处理异步任务，支持文本对话和图片识别
+    """
     try:
-        logging.info(f"开始处理异步任务: {task_id}")
+        logging.info(f"开始处理异步任务: {task_id}, 类型: {message_type}")
+        
         if task_id not in session.response_queue:
             return
             
-        result = process_long_running_task(messages)
-        
+        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
-                messages.append({"role": "assistant", "content": "图片已生成"})
             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发送'新对话'开始新的对话"
         
@@ -455,6 +670,9 @@ session_manager = SessionManager()
 
 @app.route('/api/wx', methods=['GET', 'POST'])
 def wechatai():
+    """
+    处理微信公众号请求的主要路由函数
+    """
     if request.method == 'GET':
         signature = request.args.get('signature')
         timestamp = request.args.get('timestamp')
@@ -467,153 +685,228 @@ def wechatai():
     
     try:
         encrypt_type = request.args.get('encrypt_type', '')
+        current_time = int(time.time())
         
         if encrypt_type == 'aes':
             msg_signature = request.args.get('msg_signature')
             timestamp = request.args.get('timestamp')
             nonce = request.args.get('nonce')
             
+            # 解析加密消息
             xml_tree = ET.fromstring(request.data)
             encrypted_text = xml_tree.find('Encrypt').text
             
+            # 验证消息签名
             if not verify_msg_signature(msg_signature, timestamp, nonce, TOKEN, encrypted_text):
+                logging.error("消息签名验证失败")
                 return 'Invalid signature', 403
             
-            decrypted_xml = session_manager.crypto.decrypt(encrypted_text)
-            message_data = parse_xml_message(decrypted_xml)
+            try:
+                # 解密消息
+                decrypted_xml = session_manager.crypto.decrypt(encrypted_text)
+                message_data = parse_xml_message(decrypted_xml)
+            except Exception as e:
+                logging.error(f"消息解密失败: {str(e)}")
+                return 'Decryption failed', 403
         else:
-            message_data = parse_xml_message(request.data)
+            # 处理未加密的消息
+            try:
+                message_data = parse_xml_message(request.data)
+            except Exception as e:
+                logging.error(f"解析XML消息失败: {str(e)}")
+                return 'Invalid XML', 400
         
-        user_content = message_data['content'].strip()
         from_user = message_data['from_user']
         to_user = message_data['to_user']
+        msg_type = message_data['msg_type']
         
-        logging.info(f"收到用户({from_user})消息: {user_content}")
-        session = session_manager.get_session(from_user)
-
-        if user_content == '新对话':
-            session_manager.clear_session(from_user)
+        logging.info(f"收到用户({from_user})消息类型: {msg_type}")
+        
+        # 获取或创建用户会话
+        try:
+            session = session_manager.get_session(from_user)
+        except Exception as e:
+            logging.error(f"获取用户会话失败: {str(e)}")
             return generate_response_xml(
                 from_user,
                 to_user,
-                append_status_message('已开始新的对话。请描述您的问题。'),
+                append_status_message('系统错误，请稍后重试。'),
                 encrypt_type=encrypt_type
             )
 
-        if user_content == '继续':
-            if session.pending_parts:
-                next_part = session.pending_parts.pop(0)
-                has_more = bool(session.pending_parts)
+        if msg_type == 'image':
+            try:
+                logging.info(f"收到图片消息: MediaId={message_data.get('media_id')}")
+                
+                # 创建新的异步任务
+                task_id = str(uuid.uuid4())
+                session.current_task = task_id
+                session.response_queue[task_id] = AsyncResponse()
+                
+                # 启动异步图片识别任务
+                executor.submit(
+                    handle_async_task, 
+                    session, 
+                    task_id, 
+                    None,
+                    'image',
+                    message_data
+                )
+                
                 return generate_response_xml(
                     from_user,
                     to_user,
-                    append_status_message(next_part, has_more),
+                    append_status_message("正在分析图片，请稍候...\n回复'查询'获取分析结果", is_processing=True),
                     encrypt_type=encrypt_type
                 )
-            return generate_response_xml(
-                from_user,
-                to_user,
-                append_status_message('没有更多内容了。请继续您的问题。'),
-                encrypt_type=encrypt_type
-            )
-
-        if user_content == '查询':
-            if session.current_task:
-                task_response = session.response_queue.get(session.current_task)
-                if task_response:
-                    if task_response.is_expired():
-                        del session.response_queue[session.current_task]
-                        session.current_task = None
+            except Exception as e:
+                logging.error(f"处理图片消息时发生错误: {str(e)}")
+                return generate_response_xml(
+                    from_user,
+                    to_user,
+                    append_status_message('处理图片时出现错误，请稍后重试。'),
+                    encrypt_type=encrypt_type
+                )
+        else:
+            # 处理文本消息
+            try:
+                user_content = message_data['content'].strip()
+                
+                # 处理特殊命令
+                if user_content == '新对话':
+                    session_manager.clear_session(from_user)
+                    return generate_response_xml(
+                        from_user,
+                        to_user,
+                        append_status_message('已开始新的对话。请描述您的问题。'),
+                        encrypt_type=encrypt_type
+                    )
+
+                if user_content == '继续':
+                    if session.pending_parts:
+                        next_part = session.pending_parts.pop(0)
+                        has_more = bool(session.pending_parts)
                         return generate_response_xml(
                             from_user,
                             to_user,
-                            append_status_message('请求已过期，请重新提问。'),
+                            append_status_message(next_part, has_more),
                             encrypt_type=encrypt_type
                         )
-                        
-                    if task_response.status == "completed":
-                        if task_response.response_type == "image":
-                            logging.info("返回图片响应")
-                            del session.response_queue[session.current_task]
-                            session.current_task = None
-                            return generate_response_xml(
-                                from_user,
-                                to_user,
-                                "",
-                                response_type="image",
-                                media_id=task_response.media_id,
-                                encrypt_type=encrypt_type
-                            )
-                        else:
-                            response = task_response.result
-                            del session.response_queue[session.current_task]
-                            session.current_task = None
+                    return generate_response_xml(
+                        from_user,
+                        to_user,
+                        append_status_message('没有更多内容了。请继续您的问题。'),
+                        encrypt_type=encrypt_type
+                    )
+
+                if user_content == '查询':
+                    if session.current_task:
+                        task_response = session.response_queue.get(session.current_task)
+                        if task_response:
+                            if task_response.is_expired():
+                                del session.response_queue[session.current_task]
+                                session.current_task = None
+                                return generate_response_xml(
+                                    from_user,
+                                    to_user,
+                                    append_status_message('请求已过期，请重新提问。'),
+                                    encrypt_type=encrypt_type
+                                )
                             
-                            if len(response) > 500:
-                                parts = split_message(response)
-                                first_part = parts.pop(0)
-                                session.pending_parts = parts
+                            if task_response.status == "completed":
+                                if task_response.response_type == "image":
+                                    logging.info("返回图片响应")
+                                    del session.response_queue[session.current_task]
+                                    session.current_task = None
+                                    return generate_response_xml(
+                                        from_user,
+                                        to_user,
+                                        "",
+                                        response_type="image",
+                                        media_id=task_response.media_id,
+                                        encrypt_type=encrypt_type
+                                    )
+                                else:
+                                    response = task_response.result
+                                    del session.response_queue[session.current_task]
+                                    session.current_task = None
+                                    
+                                    if len(response) > 500:
+                                        parts = split_message(response)
+                                        first_part = parts.pop(0)
+                                        session.pending_parts = parts
+                                        return generate_response_xml(
+                                            from_user,
+                                            to_user,
+                                            append_status_message(first_part, True),
+                                            encrypt_type=encrypt_type
+                                        )
+                                    return generate_response_xml(
+                                        from_user,
+                                        to_user,
+                                        append_status_message(response),
+                                        encrypt_type=encrypt_type
+                                    )
+                            elif task_response.status == "failed":
+                                error_message = '处理过程中出现错误，请重新提问。'
+                                del session.response_queue[session.current_task]
+                                session.current_task = None
                                 return generate_response_xml(
                                     from_user,
                                     to_user,
-                                    append_status_message(first_part, True),
+                                    append_status_message(error_message),
                                     encrypt_type=encrypt_type
                                 )
-                            return generate_response_xml(
-                                from_user,
-                                to_user,
-                                append_status_message(response),
-                                encrypt_type=encrypt_type
-                            )
-                    elif task_response.status == "failed":
-                        error_message = '处理过程中出现错误，请重新提问。'
-                        del session.response_queue[session.current_task]
-                        session.current_task = None
-                        return generate_response_xml(
-                            from_user,
-                            to_user,
-                            append_status_message(error_message),
-                            encrypt_type=encrypt_type
-                        )
-                    else:
-                        return generate_response_xml(
-                            from_user,
-                            to_user,
-                            append_status_message('正在处理中，请稍后再次查询。（生图需要时间）', is_processing=True),
-                            encrypt_type=encrypt_type
-                        )
-            return generate_response_xml(
-                from_user,
-                to_user,
-                append_status_message('没有正在处理的请求。'),
-                encrypt_type=encrypt_type
-            )
-
-        session.messages.append({"role": "user", "content": user_content})
-        
-        task_id = str(uuid.uuid4())
-        session.current_task = task_id
-        session.response_queue[task_id] = AsyncResponse()
-        
-        executor.submit(handle_async_task, session, task_id, session.messages.copy())
-        
-        logging.info("返回初始响应")
-        return generate_response_xml(
-            from_user,
-            to_user,
-            append_status_message(generate_initial_response(), is_processing=True),
-            encrypt_type=encrypt_type
-        )
+                            else:
+                                return generate_response_xml(
+                                    from_user,
+                                    to_user,
+                                    append_status_message('正在处理中，请稍后再次查询。', is_processing=True),
+                                    encrypt_type=encrypt_type
+                                )
+                    return generate_response_xml(
+                        from_user,
+                        to_user,
+                        append_status_message('没有正在处理的请求。'),
+                        encrypt_type=encrypt_type
+                    )
+
+                # 处理普通文本消息
+                session.messages.append({"role": "user", "content": user_content})
+                
+                # 创建新的异步任务
+                task_id = str(uuid.uuid4())
+                session.current_task = task_id
+                session.response_queue[task_id] = AsyncResponse()
+                
+                # 启动异步处理
+                executor.submit(handle_async_task, session, task_id, session.messages.copy())
+                
+                return generate_response_xml(
+                    from_user,
+                    to_user,
+                    append_status_message(generate_initial_response(), is_processing=True),
+                    encrypt_type=encrypt_type
+                )
+                
+            except Exception as e:
+                logging.error(f"处理文本消息时发生错误: {str(e)}")
+                return generate_response_xml(
+                    from_user,
+                    to_user,
+                    append_status_message('处理消息时出现错误，请稍后重试。'),
+                    encrypt_type=encrypt_type
+                )
 
     except Exception as e:
         logging.error(f"处理请求时出错: {str(e)}")
         return generate_response_xml(
-            message_data['from_user'],
-            message_data['to_user'],
+            message_data['from_user'] if 'message_data' in locals() else 'unknown',
+            message_data['to_user'] if 'message_data' in locals() else 'unknown',
             append_status_message('抱歉，系统暂时出现问题，请稍后重试。'),
             encrypt_type if 'encrypt_type' in locals() else ''
         )
-
+    
 def cleanup_sessions():
     while True:
         time.sleep(3600)  # 每小时清理一次
@@ -626,4 +919,4 @@ if __name__ == '__main__':
     cleanup_thread = threading.Thread(target=cleanup_sessions, daemon=True)
     cleanup_thread.start()
     
-    app.run(host='0.0.0.0', port=7860, debug=True)
+    app.run(host='0.0.0.0', port=7860, debug=True)