src/core/message_processor.py

# src/core/message_processor.py
import json
import time
import base64
import uuid
import requests
from curl_cffi import requests as curl_requests

from src.core.logger import logger # 从新的位置导入 logger
from src.core.utils import Utils # 从新的位置导入 Utils
from app import CONFIG, DEFAULT_HEADERS # 导入 CONFIG 和 DEFAULT_HEADERS

class MessageProcessor:
    @staticmethod
    def create_chat_response(message, model, is_stream=False):
        base_response = {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "created": int(time.time()),
            "model": model
        }
        if is_stream:
            return {
                **base_response,
                "object": "chat.completion.chunk",
                "choices": [{
                    "index": 0,
                    "delta": {
                        "content": message
                    }
                }]
            }
        return {
            **base_response,
            "object": "chat.completion",
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": message
                },
                "finish_reason": "stop"
            }],
            "usage": None
        }

    def process_model_response(response, model):
        result = {"token": None, "imageUrl": None}
        if CONFIG["IS_IMG_GEN"]:
            if response.get("cachedImageGenerationResponse") and not CONFIG["IS_IMG_GEN2"]:
                result["imageUrl"] = response["cachedImageGenerationResponse"]["imageUrl"]
                return result

        if model == 'grok-2':
            result["token"] = response.get("token")
        elif model in ['grok-2-search', 'grok-3-search']:
            if response.get("webSearchResults") and CONFIG["ISSHOW_SEARCH_RESULTS"]:
                result["token"] = f"\r\n{Utils.organize_search_results(response['webSearchResults'])}\r\n"
            else:
                result["token"] = response.get("token")
        elif model == 'grok-3':
            result["token"] = response.get("token")
        elif model == 'grok-3-deepsearch':
            if response.get("messageStepId") and not CONFIG["SHOW_THINKING"]:
                return result
            if response.get("messageStepId") and not CONFIG["IS_THINKING"]:
                result["token"] = "" + response.get("token", "")
                CONFIG["IS_THINKING"] = True
            elif not response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "final":
                result["token"] = "" + response.get("token", "")
                CONFIG["IS_THINKING"] = False
            elif (response.get("messageStepId") and CONFIG["IS_THINKING"] and response.get("messageTag") == "assistant") or response.get("messageTag") == "final":
                result["token"] = response.get("token")
        elif model == 'grok-3-reasoning':
            if response.get("isThinking") and not CONFIG["SHOW_THINKING"]:
                return result
            if response.get("isThinking") and not CONFIG["IS_THINKING"]:
                result["token"] = "" + response.get("token", "")
                CONFIG["IS_THINKING"] = True
            elif not response.get("isThinking") and CONFIG["IS_THINKING"]:
                result["token"] = "" + response.get("token", "")
                CONFIG["IS_THINKING"] = False
            else:
                result["token"] = response.get("token")

        return result

    @staticmethod # 将方法改为静态方法
    def handle_image_response(image_url):
        max_retries = 2
        retry_count = 0
        image_base64_response = None
        while retry_count < max_retries:
            try:
                proxy_options = Utils.get_proxy_options()
                image_base64_response = curl_requests.get(
                    f"https://assets.grok.com/{image_url}",
                    headers={
                        **DEFAULT_HEADERS,
                        "Cookie":CONFIG["SERVER"]['COOKIE']
                    },
                    impersonate="chrome133a",
                    **proxy_options
                )
                if image_base64_response.status_code == 200:
                    break
                retry_count += 1
                if retry_count == max_retries:
                    raise Exception(f"上游服务请求失败! status: {image_base64_response.status_code}")
                time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)
            except Exception as error:
                logger.error(str(error), "Server")
                retry_count += 1
                if retry_count == max_retries:
                    raise
                time.sleep(CONFIG["API"]["RETRY_TIME"] / 1000 * retry_count)

        image_buffer = image_base64_response.content

        if not CONFIG["API"]["PICGO_KEY"] and not CONFIG["API"]["TUMY_KEY"]:
            base64_image = base64.b64encode(image_buffer).decode('utf-8')
            image_content_type = image_base64_response.headers.get('content-type', 'image/jpeg')
            return f"![image](data:{image_content_type};base64,{base64_image})"

        logger.info("开始上传图床", "Server")
        if CONFIG["API"]["PICGO_KEY"]:
            files = {'source': ('image.jpg', image_buffer, 'image/jpeg')}
            headers = {
                "X-API-Key": CONFIG["API"]["PICGO_KEY"]
            }
            response_url = requests.post(
                "https://www.picgo.net/api/1/upload",
                files=files,
                headers=headers
            )
            if response_url.status_code != 200:
                return "生图失败，请查看PICGO图床密钥是否设置正确"
            else:
                logger.info("生图成功", "Server")
                result = response_url.json()
                return f"![image]({result['image']['url']})"
        elif CONFIG["API"]["TUMY_KEY"]:
            files = {'file': ('image.jpg', image_buffer, 'image/jpeg')}
            headers = {
                "Accept": "application/json",
                'Authorization': f"Bearer {CONFIG['API']['TUMY_KEY']}"
            }
            response_url = requests.post(
                "https://tu.my/api/v1/upload",
                files=files,
                headers=headers
            )
            if response_url.status_code != 200:
                return "生图失败，请查看TUMY图床密钥是否设置正确"
            else:
                try:
                    result = response_url.json()
                    logger.info("生图成功", "Server")
                    return f"![image]({result['data']['links']['url']})"
                except Exception as error:
                    logger.error(str(error), "Server")
                    return "生图失败，请查看TUMY图床密钥是否设置正确"


    def handle_non_stream_response(response, model):
        try:
            logger.info("开始处理非流式响应", "Server")
            stream = response.iter_lines()
            full_response = ""
            CONFIG["IS_THINKING"] = False
            CONFIG["IS_IMG_GEN"] = False
            CONFIG["IS_IMG_GEN2"] = False
            for chunk in stream:
                if not chunk:
                    continue
                try:
                    line_json = json.loads(chunk.decode("utf-8").strip())
                    if line_json.get("error"):
                        logger.error(json.dumps(line_json, indent=2), "Server")
                        yield json.dumps({"error": "RateLimitError"}) + "\n\n"
                        return

                    response_data = line_json.get("result", {}).get("response")
                    if not response_data:
                        continue

                    if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
                        CONFIG["IS_IMG_GEN"] = True

                    result = MessageProcessor.process_model_response(response_data, model) # 使用类名调用静态方法
                    if result["token"]:
                        full_response += result["token"]
                    if result["imageUrl"]:
                        CONFIG["IS_IMG_GEN2"] = True
                        return MessageProcessor.handle_image_response(result["imageUrl"]) # 使用类名调用静态方法

                except json.JSONDecodeError:
                    continue
                except Exception as e:
                    logger.error(f"处理流式响应行时出错: {str(e)}", "Server")
                    continue
            return full_response
        except Exception as error:
            logger.error(str(error), "Server")
            raise

    def handle_stream_response(response, model):
        def generate():
            logger.info("开始处理流式响应", "Server")
            stream = response.iter_lines()
            CONFIG["IS_THINKING"] = False
            CONFIG["IS_IMG_GEN"] = False
            CONFIG["IS_IMG_GEN2"] = False
            for chunk in stream:
                if not chunk:
                    continue
                try:
                    line_json = json.loads(chunk.decode("utf-8").strip())
                    if line_json.get("error"):
                        logger.error(json.dumps(line_json, indent=2), "Server")
                        yield json.dumps({"error": "RateLimitError"}) + "\n\n"
                        return

                    response_data = line_json.get("result", {}).get("response")
                    if not response_data:
                        continue

                    if response_data.get("doImgGen") or response_data.get("imageAttachmentInfo"):
                        CONFIG["IS_IMG_GEN"] = True

                    result = MessageProcessor.process_model_response(response_data, model) # 使用类名调用静态方法
                    if result["token"]:
                        yield f"data: {json.dumps(MessageProcessor.create_chat_response(result['token'], model, True))}\n\n"
                    if result["imageUrl"]:
                        CONFIG["IS_IMG_GEN2"] = True
                        image_data = MessageProcessor.handle_image_response(result["imageUrl"]) # 使用类名调用静态方法
                        yield f"data: {json.dumps(MessageProcessor.create_chat_response(image_data, model, True))}\n\n"

                except json.JSONDecodeError:
                    continue
                except Exception as e:
                    logger.error(f"处理流式响应行时出错: {str(e)}", "Server")
                    continue

            yield "data: [DONE]\n\n"

        return generate()