Update app.py

app.py

@@ -1,134 +1,84 @@
 import os
-from flask import Flask, request, abort
-import hashlib
-import time
-import xmltodict
+from flask import Flask, request, make_response
+import xml.etree.ElementTree as ET
+from wechatpy import parse_message, create_reply
+from wechatpy.utils import check_signature
+from wechatpy.exceptions import InvalidSignatureException
 from openai import OpenAI
-import re
-import base64
-import requests
 
 app = Flask(__name__)
 
-# 环境变量配置
+# 配置
 TOKEN = os.getenv("TOKEN")
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-APPID = os.getenv("APPID")
-APPSECRET = os.getenv("APPSECRET")
+API_KEY = os.getenv("OPENAI_API_KEY")
+BASE_URL = os.getenv("OPENAI_BASE_URL")
+client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
 
-client = OpenAI(api_key=OPENAI_API_KEY)
+# 存储用户会话信息
+user_sessions = {}
 
-# 用户模型存储
-user_models = {}
-
-def check_signature():
-    signature = request.args.get('signature', '')
-    timestamp = request.args.get('timestamp', '')
-    nonce = request.args.get('nonce', '')
-    token = TOKEN
-    tmp_list = sorted([token, timestamp, nonce])
-    tmp_str = ''.join(tmp_list)
-    hash_obj = hashlib.sha1(tmp_str.encode('utf-8'))
-    return hash_obj.hexdigest() == signature
-
-def split_message(message, max_length=500):
-    """Split a message into chunks of max_length characters."""
-    return [message[i:i+max_length] for i in range(0, len(message), max_length)]
-
-def get_openai_response(message, model="gpt-4o", image_url=None):
+def get_openai_response(messages, model="gpt-3.5-turbo"):
     try:
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": message}
-        ]
-        
-        if image_url:
-            messages[1]["content"] = [
-                {"type": "text", "text": message},
-                {"type": "image_url", "image_url": {"url": image_url}}
-            ]
-        
-        completion = client.chat.completions.create(
+        response = client.chat.completions.create(
             model=model,
-            messages=messages,
-            max_tokens=300
+            messages=messages
         )
-        return completion.choices[0].message.content
+        return response.choices[0].message.content
     except Exception as e:
-        return f"Error: {str(e)}"
-
-def handle_text_message(from_user, to_user, content):
-    if content.startswith('/'):
-        # 处理命令
-        parts = content.split()
-        command = parts[0][1:]
-        if command == 'help':
-            return "Available commands: /help, /setmodel [model_name]"
-        elif command == 'setmodel':
-            if len(parts) > 1:
-                model = parts[1]
-                user_models[from_user] = model
-                return f"Model switched to {model}"
-            else:
-                return "Please specify a model name. Usage: /setmodel [model_name]"
-        else:
-            return "Unknown command. Type /help for available commands."
-    
-    # 正常对话，调用OpenAI API
-    model = user_models.get(from_user, "gpt-4o")
-    response = get_openai_response(content, model)
-    return format_reply(from_user, to_user, response)
-
-def handle_image_message(from_user, to_user, pic_url):
-    model = user_models.get(from_user, "gpt-4o")
-    response = get_openai_response("What's in this image?", model, pic_url)
-    return format_reply(from_user, to_user, response)
+        print(f"调用OpenAI API时出错: {str(e)}")
+        return "抱歉，我遇到了一些问题，无法回答您的问题。"
 
-def format_reply(from_user, to_user, content):
-    response_parts = split_message(content)
-    replies = []
-    for part in response_parts:
-        reply = f"""
-        <xml>
-        <ToUserName><![CDATA[{from_user}]]></ToUserName>
-        <FromUserName><![CDATA[{to_user}]]></FromUserName>
-        <CreateTime>{int(time.time())}</CreateTime>
-        <MsgType><![CDATA[text]]></MsgType>
-        <Content><![CDATA[{part}]]></Content>
-        </xml>
-        """
-        replies.append(reply)
-    return ''.join(replies)
+def split_message(message, max_length=500):
+    return [message[i:i+max_length] for i in range(0, len(message), max_length)]
 
 @app.route('/', methods=['GET', 'POST'])
 def wechat():
     if request.method == 'GET':
-        echostr = request.args.get('echostr', '')
-        if check_signature():
-            return echostr
-        abort(403)
-    elif request.method == 'POST':
-        if not check_signature():
-            abort(403)
-        
-        xml_data = request.data
-        msg = xmltodict.parse(xml_data)['xml']
-        msg_type = msg['MsgType']
-        from_user = msg['FromUserName']
-        to_user = msg['ToUserName']
+        token = TOKEN
+        query = request.args
+        signature = query.get('signature', '')
+        timestamp = query.get('timestamp', '')
+        nonce = query.get('nonce', '')
+        echostr = query.get('echostr', '')
+        try:
+            check_signature(token, signature, timestamp, nonce)
+        except InvalidSignatureException:
+            return 'Invalid signature'
+        return echostr
 
-        if msg_type == 'text':
-            content = msg['Content']
-            return handle_text_message(from_user, to_user, content)
-        elif msg_type == 'image':
-            pic_url = msg['PicUrl']
-            return handle_image_message(from_user, to_user, pic_url)
-        elif msg_type == 'event':
-            event = msg['Event']
-            if event == 'subscribe':
-                return format_reply(from_user, to_user, "感谢关注！输入 /help 查看可用命令。")
+    if request.method == 'POST':
+        xml_str = request.data
+        if not xml_str:
+            return ""
         
-        return 'success'
+        msg = parse_message(xml_str)
+        if msg.type == 'text':
+            user_id = msg.source
+            content = msg.content
+
+            if content.startswith('/model'):
+                # 切换模型
+                model = content.split(' ')[1]
+                user_sessions[user_id] = {'model': model, 'messages': []}
+                return create_reply(f'模型已切换为 {model}', msg).render()
+
+            if user_id not in user_sessions:
+                user_sessions[user_id] = {'model': 'gpt-3.5-turbo', 'messages': []}
+
+            session = user_sessions[user_id]
+            session['messages'].append({"role": "user", "content": content})
+            
+            response = get_openai_response(session['messages'], session['model'])
+            session['messages'].append({"role": "assistant", "content": response})
+
+            # 分割长消息
+            response_parts = split_message(response)
+            for part in response_parts[:-1]:
+                create_reply(part, msg).render()
+            
+            return create_reply(response_parts[-1], msg).render()
+
+        return create_reply('Sorry, can not handle this for now', msg).render()
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)