Update app.py

app.py

@@ -6,10 +6,12 @@ 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
@@ -19,7 +21,6 @@ import struct
 import random
 import string
 import requests
-from typing import Optional, Dict, Any
 
 logging.basicConfig(
     level=logging.INFO,
@@ -46,13 +47,13 @@ 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 the AI model
+# Define tools for image generation
 TOOLS = [
     {
         "type": "function",
         "function": {
             "name": "generate_image",
-            "description": "Generate an image based on text description and return its markdown URL",
+            "description": "Generate an image based on text description",
             "parameters": {
                 "type": "object",
                 "properties": {
@@ -105,8 +106,8 @@ class AsyncResponse:
         self.error = None
         self.create_time = time.time()
         self.timeout = 3600
-        self.media_id = None
-        self.response_type = "text"
+        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
@@ -170,71 +171,6 @@ class SessionManager:
                 del self.sessions[user_id]
                 logging.info(f"已清理过期会话: {user_id}")
 
-class ImageService:
-    @staticmethod
-    def generate_image(prompt: str) -> str:
-        try:
-            logging.info(f"开始生成图片，提示词: {prompt}")
-            
-            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": prompt
-                    }],
-                    "stream": False
-                }
-            )
-            
-            logging.info(f"图片生成服务响应状态码: {response.status_code}")
-            response.raise_for_status()
-            
-            result = response.json()
-            logging.info(f"图片生成服务响应内容: {json.dumps(result, ensure_ascii=False)}")
-            
-            if not result.get('choices') or not result['choices'][0].get('message', {}).get('content'):
-                raise ValueError("Invalid response format")
-                
-            image_url = result['choices'][0]['message']['content']
-            logging.info(f"成功获取图片URL: {image_url}")
-            
-            return image_url
-        except Exception as e:
-            logging.error(f"Image generation error: {str(e)}")
-            raise
-
-    @staticmethod
-    def get_media_id(image_url: str) -> str:
-        try:
-            logging.info(f"开始下载图片: {image_url}")
-            image_response = requests.get(image_url)
-            image_response.raise_for_status()
-            image_data = image_response.content
-            
-            logging.info("开始上传图片到微信服务器")
-            upload_url = f'https://api.weixin.qq.com/cgi-bin/media/upload?access_token={TOKEN}&type=image'
-            files = {'media': ('image.jpg', image_data, 'image/jpeg')}
-            response = requests.post(upload_url, files=files)
-            response.raise_for_status()
-            
-            result = response.json()
-            logging.info(f"微信服务器响应: {json.dumps(result, ensure_ascii=False)}")
-            
-            if 'media_id' not in result:
-                raise ValueError(f"Failed to get media_id: {result}")
-                
-            logging.info(f"成功获取media_id: {result['media_id']}")
-            return result['media_id']
-        except Exception as e:
-            logging.error(f"WeChat media upload error: {str(e)}")
-            raise
-
 def convert_markdown_to_wechat(md_text):
     if not md_text:
         return md_text
@@ -285,7 +221,6 @@ def generate_response_xml(to_user, from_user, content, response_type='text', med
     nonce = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
     
     if response_type == 'image' and media_id:
-        logging.info(f"生成图片消息响应，media_id: {media_id}")
         xml_content = f'''
         <xml>
             <ToUserName><![CDATA[{to_user}]]></ToUserName>
@@ -299,7 +234,6 @@ def generate_response_xml(to_user, from_user, content, response_type='text', med
         '''
     else:
         formatted_content = convert_markdown_to_wechat(content)
-        logging.info(f"生成文本消息响应: {formatted_content}")
         xml_content = f'''
         <xml>
             <ToUserName><![CDATA[{to_user}]]></ToUserName>
@@ -331,81 +265,69 @@ def generate_response_xml(to_user, from_user, content, response_type='text', med
     response.content_type = 'application/xml'
     return response
 
-def generate_initial_response():
-    return "您的请求正在处理中，请回复'查询'获取结果"
-
-def split_message(message, max_length=500):
-    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):
-    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
-
-def process_ai_response(messages):
+def process_long_running_task(messages):
     try:
-        logging.info("开始处理AI响应")
-        completion = client.chat.completions.create(
+        logging.info("开始调用AI服务")
+        response = client.chat.completions.create(
             model="o3-mini",
             messages=messages,
             tools=TOOLS,
-            tool_choice="auto"
+            tool_choice="auto",
+            timeout=60
         )
-        logging.info("收到AI响应")
-        
-        # Handle tool calls if present
-        if completion.choices[0].message.tool_calls:
-            logging.info("检测到工具调用")
-            for tool_call in completion.choices[0].message.tool_calls:
-                if tool_call.function.name == "generate_image":
-                    try:
-                        logging.info("开始执行图片生成")
-                        args = json.loads(tool_call.function.arguments)
-                        # Generate image and get markdown URL
-                        image_url = ImageService.generate_image(args['prompt'])
-                        
-                        # Get WeChat media_id
-                        media_id = ImageService.get_media_id(image_url)
-                        
-                        messages.append({
-                            "role": "assistant",
-                            "content": f"已生成图片"
-                        })
-                        
-                        return {
-                            "type": "image",
-                            "content": None,
-                            "media_id": media_id
-                        }
-                    except Exception as e:
-                        logging.error(f"图片生成过程失败: {str(e)}")
-                        return {
-                            "type": "text",
-                            "content": f"抱歉，图片生成失败：{str(e)}",
-                            "media_id": None
-                        }
-        
-        # Handle normal text response
-        logging.info("处理普通文本响应")
-        response_content = completion.choices[0].message.content
-        messages.append({
-            "role": "assistant",
-            "content": response_content
-        })
+        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']
+                        }]
+                    }
+                )
+                image_response.raise_for_status()
+                result = image_response.json()
+                logging.info("图片生成成功，准备下载图片")
+
+                image_url = result['choices'][0]['message']['content']
+                img_response = requests.get(image_url)
+                img_response.raise_for_status()
+                
+                logging.info("开始上传图片到微信服务器")
+                upload_url = f'https://api.weixin.qq.com/cgi-bin/media/upload?access_token={TOKEN}&type=image'
+                files = {'media': ('image.jpg', img_response.content, 'image/jpeg')}
+                upload_response = requests.post(upload_url, files=files)
+                upload_response.raise_for_status()
+                media_result = upload_response.json()
+                
+                if 'media_id' not in media_result:
+                    raise ValueError(f"Failed to get media_id: {media_result}")
+                
+                logging.info(f"图片上传成功，media_id: {media_result['media_id']}")
+                return {
+                    "type": "image",
+                    "media_id": media_result['media_id']
+                }
+
+        logging.info("返回文本响应")
         return {
             "type": "text",
-            "content": response_content,
-            "media_id": None
+            "content": response.choices[0].message.content
         }
     except Exception as e:
-        logging.error(f"AI处理错误: {str(e)}")
+        logging.error(f"API调用错误: {str(e)}")
         raise
 
 def handle_async_task(session, task_id, messages):
@@ -414,20 +336,42 @@ def handle_async_task(session, task_id, messages):
         if task_id not in session.response_queue:
             return
             
-        result = process_ai_response(messages)
-        logging.info(f"异步任务处理完成: {task_id}")
+        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].result = result.get("content")
-            session.response_queue[task_id].response_type = result.get("type")
-            session.response_queue[task_id].media_id = result.get("media_id")
+            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"]
+                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):
+    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):
+    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()
 
 @app.route('/api/wx', methods=['GET', 'POST'])
@@ -509,8 +453,10 @@ def wechatai():
                         )
                         
                     if task_response.status == "completed":
-                        response_type = task_response.response_type
-                        if response_type == "image":
+                        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,
@@ -564,27 +510,21 @@ def wechatai():
                 encrypt_type=encrypt_type
             )
 
-        # Regular message processing
-        logging.info("准备开始处理用户消息")
         session.messages.append({"role": "user", "content": user_content})
         
         task_id = str(uuid.uuid4())
         session.current_task = task_id
         session.response_queue[task_id] = AsyncResponse()
         
-        # Submit task to executor
-        logging.info(f"提交任务到执行器: {task_id}")
         executor.submit(handle_async_task, session, task_id, session.messages.copy())
         
-        # Return immediate response
-        initial_response = generate_response_xml(
+        logging.info("返回初始响应")
+        return generate_response_xml(
             from_user,
             to_user,
             append_status_message(generate_initial_response(), is_processing=True),
             encrypt_type=encrypt_type
         )
-        logging.info("返回初始响应给用户")
-        return initial_response
 
     except Exception as e:
         logging.error(f"处理请求时出错: {str(e)}")