ai-on-wechat/plugins/file4upload/file4upload0.txt

import requests
import json
import re
import plugins
from bridge.reply import Reply, ReplyType
from bridge.context import ContextType
from channel.chat_message import ChatMessage
from plugins import *
from common.log import logger
from common.expired_dict import ExpiredDict
import os
from docx import Document
import markdown
import fitz
from openpyxl import load_workbook
import csv
from bs4 import BeautifulSoup
from pptx import Presentation
from PIL import Image
import base64
import html

EXTENSION_TO_TYPE = {
    'pdf': 'pdf',
    'doc': 'docx', 'docx': 'docx',
    'md': 'md',
    'txt': 'txt',
    'xls': 'excel', 'xlsx': 'excel',
    'csv': 'csv',
    'html': 'html', 'htm': 'html',
    'ppt': 'ppt', 'pptx': 'ppt'
}

@plugins.register(
    name="file4upload",
    desire_priority=-1,
    desc="A plugin for upload",
    version="0.0.01",
    author="",
)
class file4upload(Plugin):
    def __init__(self):
        super().__init__()
        try:
            curdir = os.path.dirname(__file__)
            config_path = os.path.join(curdir, "config.json")
            if os.path.exists(config_path):
                with open(config_path, "r", encoding="utf-8") as f:
                    self.config = json.load(f)
            else:
                # 使用父类的方法来加载配置
                self.config = super().load_config()

                if not self.config:
                    raise Exception("config.json not found")
            # 设置事件处理函数
            self.handlers[Event.ON_HANDLE_CONTEXT] = self.on_handle_context
            self.params_cache = ExpiredDict(300)

             # 从配置中提取所需的设置
            self.keys = self.config.get("keys", {})
            self.url_sum = self.config.get("url_sum", {})
            self.search_sum = self.config.get("search_sum", {})
            self.file_sum = self.config.get("file_sum", {})
            self.image_sum = self.config.get("image_sum", {})
            self.note = self.config.get("note", {})

            self.sum4all_key = self.keys.get("sum4all_key", "")
            self.search1api_key = self.keys.get("search1api_key", "")
            self.gemini_key = self.keys.get("gemini_key", "")
            self.bibigpt_key = self.keys.get("bibigpt_key", "")
            self.outputLanguage = self.keys.get("outputLanguage", "zh-CN")
            self.opensum_key = self.keys.get("opensum_key", "")
            self.open_ai_api_key = self.keys.get("open_ai_api_key", "")
            self.model = self.keys.get("model", "gpt-3.5-turbo")
            self.open_ai_api_base = self.keys.get("open_ai_api_base", "https://api.openai.com/v1")
            self.xunfei_app_id = self.keys.get("xunfei_app_id", "")
            self.xunfei_api_key = self.keys.get("xunfei_api_key", "")
            self.xunfei_api_secret = self.keys.get("xunfei_api_secret", "")
            self.perplexity_key = self.keys.get("perplexity_key", "")
            self.flomo_key = self.keys.get("flomo_key", "")
            # 之前提示
            self.previous_prompt=''

            self.file_sum_enabled = self.file_sum.get("enabled", False)
            self.file_sum_service = self.file_sum.get("service", "")
            self.max_file_size = self.file_sum.get("max_file_size", 15000)
            self.file_sum_group = self.file_sum.get("group", True)
            self.file_sum_qa_prefix = self.file_sum.get("qa_prefix", "问")
            self.file_sum_prompt = self.file_sum.get("prompt", "")

            self.image_sum_enabled = self.image_sum.get("enabled", False)
            self.image_sum_service = self.image_sum.get("service", "")
            self.image_sum_group = self.image_sum.get("group", True)
            self.image_sum_qa_prefix = self.image_sum.get("qa_prefix", "问")
            self.image_sum_prompt = self.image_sum.get("prompt", "")

            # 初始化成功日志
            logger.info("[file4upload] inited.")
        except Exception as e:
            # 初始化失败日志
            logger.warn(f"file4upload init failed: {e}")

    def on_handle_context(self, e_context: EventContext):
        context = e_context["context"]
        if context.type not in [ContextType.TEXT, ContextType.SHARING,ContextType.FILE,ContextType.IMAGE]:
            return
        msg: ChatMessage = e_context["context"]["msg"]
        user_id = msg.from_user_id
        content = context.content
        isgroup = e_context["context"].get("isgroup", False)
        logger.info(f"user_id:{user_id},content:{content},isgroup:{isgroup}")
        logger.info(f'上下文参数缓存键字典：{self.params_cache.keys}')
        logger.info(f'user_id in self.params_cache: {user_id in self.params_cache}')
        # logger.info(f'last_file_content in self.params_cache[user_id] :{'last_file_content' in self.params_cache[user_id]}')
                # 上次提示
        if  context.type == ContextType.TEXT:
            self.previous_prompt=msg.content

        if context.type == ContextType.FILE:
            logger.info('开始首次处理文件')
            # logger.info(f'{self.params_cache}')
            # # content = "The World"
            # logger.info(f'file4upload 插件作画')
            # e_context.action = EventAction.CONTINUE  # 事件继续，交付给下个插件或默认逻辑

            # e_context["context"].type = ContextType.TEXT
            # e_context["context"].content = "The World"
            # logger.info(f'file4upload 插件处理')
            # e_context.action = EventAction.CONTINUE  # 事件继续，交付给下个插件或默认逻辑

            if isgroup and not self.file_sum_group:
                # 群聊中忽略处理文件
                logger.info("群聊消息，文件处理功能已禁用")
                return
            logger.info("on_handle_context: 处理上下文开始")
            context.get("msg").prepare()
            file_path = context.content
            logger.info(f"on_handle_context: 获取到文件路径 {file_path}")
            
            # 检查是否应该进行文件总结
            if self.file_sum_enabled:
                # 更新params_cache中的last_file_content
                self.params_cache[user_id] = {}
                file_content = self.extract_content(file_path)
                if file_content is None:
                    logger.info("文件内容无法提取，跳过处理")
                else:
                    file_content=f'Use the following context as your learned knowledge, inside <context></context> XML tags.\n\t<context>{file_content}</context>\n\t\n\tWhen answer to user:\n\t- If you don\'t know, just say that you don\'t know.\n\t- If you don\'t know when you are not sure, ask for clarification.\n\tAvoid mentioning that you obtained the information from the context.\n\tAnd answer according to the language of the user\'s question.\n\t\t\t\n\tGiven the context information, answer the query.\n\tQuery:'
                    self.params_cache[user_id]['last_file_content'] = file_content
                    logger.info(f'写入文件内容缓存')
                    # logger.info('Updated last_file_content in params_cache for user.')
                    # # self.handle_file(file_content, e_context)

                    # e_context["context"].type = ContextType.TEXT
                    # e_context["context"].content = file_content+'\n\t请总结分析这份数据文件。文件内容是HRV分析报告，请分析这个用户的心梗、脑梗的风险等级（如无风险、低风险、中风险、高风险），并提供健康建议'
                    # logger.info(f'file4upload 插件处理上传文件')
                    e_context.action = EventAction.CONTINUE  # 事件继续，交付给下个插件或默认逻辑                    
            else:
                logger.info("文件总结功能已禁用，不对文件内容进行处理")

        elif context.type == ContextType.IMAGE:
            logger.info('开始首次处理图片')

        if user_id in self.params_cache and ('last_file_content' in self.params_cache[user_id] or 'last_image_base64' in self.params_cache[user_id] ):
            # logger.info('上传过文件或图片')
            if 'last_file_content' in self.params_cache[user_id]:
                logger.info('上次文件内容开始')
                e_context["context"].type = ContextType.TEXT
                e_context["context"].content = self.params_cache[user_id]['last_file_content']+f'\n\t{self.previous_prompt}'
                logger.info(f'file4upload 插件处理上传文件')
                self.previous_prompt=''
                e_context.action = EventAction.CONTINUE  # 事件继续，交付给下个插件或默认逻辑                 

            
            elif 'last_image_base64' in self.params_cache[user_id]:
                logger.info('上次图片开始')

        

    def handle_file(self, content, e_context):
        logger.info("handle_file: 向LLM发送内容总结请求")
        # 根据sum_service的值选择API密钥和基础URL
        if self.file_sum_service == "openai":
            api_key = self.open_ai_api_key
            api_base = self.open_ai_api_base
            model = self.model
        elif self.file_sum_service == "sum4all":
            api_key = self.sum4all_key
            api_base = "https://pro.sum4all.site/v1"
            model = "sum4all"
        elif self.file_sum_service == "gemini":
            api_key = self.gemini_key
            model = "gemini"
            api_base = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"
        else:
            logger.error(f"未知的sum_service配置: {self.file_sum_service}")
            return
        msg: ChatMessage = e_context["context"]["msg"]
        user_id = msg.from_user_id
        user_params = self.params_cache.get(user_id, {})
        prompt = user_params.get('prompt', self.file_sum_prompt)
        if model == "gemini":
            headers = {
                'Content-Type': 'application/json',
                'x-goog-api-key': api_key
            }
            data = {
            "contents": [
                {"role": "user", "parts": [{"text": prompt}]},
                {"role": "model", "parts": [{"text": "okay"}]},
                {"role": "user", "parts": [{"text": content}]}
            ],
            "generationConfig": {
                "maxOutputTokens": 800
            }
            }
            api_url = api_base
        else:
            headers = {
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {api_key}'
            }
            # content=f'Use the following context as your learned knowledge, inside <context></context> XML tags.\n\t<context>{content}</context>\n\t\n\tWhen answer to user:\n\t- If you don\'t know, just say that you don\'t know.\n\t- If you don\'t know when you are not sure, ask for clarification.\n\tAvoid mentioning that you obtained the information from the context.\n\tAnd answer according to the language of the user\'s question.\n\t\t\t\n\tGiven the context information, answer the query'
            if self.previous_prompt!='':
                prompt=self.previous_prompt
                logger.info(f"改变提示，使用上次提示")
            data = {
                "model": model,
                "messages": [
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": content}
                ]
            }
            api_url = f"{api_base}/chat/completions"
        try:
            logger.info(f'handle_file: 请求文件内容{json.dumps(data, ensure_ascii=False)}')
            response = requests.post(api_url, headers=headers, data=json.dumps(data))
            response.raise_for_status()
            response_data = response.json()
            
            # 解析 JSON 并获取 content
            if model == "gemini":
                if "candidates" in response_data and len(response_data["candidates"]) > 0:
                    first_candidate = response_data["candidates"][0]
                    if "content" in first_candidate:
                        if "parts" in first_candidate["content"] and len(first_candidate["content"]["parts"]) > 0:
                            response_content = first_candidate["content"]["parts"][0]["text"].strip()  # 获取响应内容
                            logger.info(f"Gemini API response content: {response_content}")  # 记录响应内容
                            reply_content = response_content.replace("\\n", "\n")  # 替换 \\n 为 \n
                        else:
                            logger.error("Parts not found in the Gemini API response content")
                            reply_content = "Parts not found in the Gemini API response content"
                    else:
                        logger.error("Content not found in the Gemini API response candidate")
                        reply_content = "Content not found in the Gemini API response candidate"
                else:
                    logger.error("No candidates available in the Gemini API response")
                    reply_content = "No candidates available in the Gemini API response"        
            else:
                if "choices" in response_data and len(response_data["choices"]) > 0:
                    first_choice = response_data["choices"][0]
                    if "message" in first_choice and "content" in first_choice["message"]:
                        response_content = first_choice["message"]["content"].strip()  # 获取响应内容
                        logger.info(f"LLM API response content")  # 记录响应内容
                        reply_content = response_content.replace("\\n", "\n")  # 替换 \\n 为 \n
                        if msg.ctype == ContextType.FILE and self.previous_prompt =='':
                            reply_content="您刚刚上传了一个文件，请问我有什么可以帮您的呢？"
                            self.previous_prompt=''
                        # 已上传过，重置 previous_prompt
                        elif self.previous_prompt !='' and 'last_file_content' in self.params_cache[user_id]:
                            logger.info(f'last_file_content 已经回答最后的提示,清空previous_prompt')
                            self.previous_prompt =''
                    else:
                        logger.error("Content not found in the response")
                        reply_content = "Content not found in the LLM API response"
                else:
                    logger.error("No choices available in the response")
                    reply_content = "No choices available in the LLM API response"

        except requests.exceptions.RequestException as e:
            logger.error(f"Error calling LLM API: {e}")
            reply_content = f"An error occurred while calling LLM API"

        reply = Reply()
        reply.type = ReplyType.TEXT
        # reply.content = f"{remove_markdown(reply_content)}\n\n💬5min内输入{self.file_sum_qa_prefix}+问题，可继续追问" 
        reply.content = f"{remove_markdown(reply_content)}"
        e_context["reply"] = reply
        e_context.action = EventAction.BREAK_PASS
    def handle_image(self, base64_image, e_context):
        logger.info("handle_image: 解析图像处理API的响应")
        msg: ChatMessage = e_context["context"]["msg"]
        user_id = msg.from_user_id
        user_params = self.params_cache.get(user_id, {})
        prompt = user_params.get('prompt', self.image_sum_prompt)

        if self.image_sum_service == "openai":
            api_key = self.open_ai_api_key
            api_base = f"{self.open_ai_api_base}/chat/completions"
            model = "gpt-4o-mini"
        elif self.image_sum_service == "xunfei":
            api_key = self.xunfei_api_key
            api_base = "https://spark.sum4all.site/v1/chat/completions"
            model = "spark-chat-vision"
        elif self.image_sum_service == "sum4all":
            api_key = self.sum4all_key
            api_base = "https://pro.sum4all.site/v1/chat/completions"
            model = "sum4all-vision"
        elif self.image_sum_service == "gemini":
            api_key = self.gemini_key
            api_base = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"
            payload = {
                "contents": [
                    {
                        "parts": [
                            {"text": prompt},
                            {
                                "inline_data": {
                                    "mime_type":"image/png",
                                    "data": base64_image
                                }
                            }
                        ]
                    }
                ]
            }
            headers = {
                "Content-Type": "application/json",
                "x-goog-api-key": api_key
            }
            logger.info(f"准备发送请求. Payload大小: {len(json.dumps(payload))} 字节")

        else:
            logger.error(f"未知的image_sum_service配置: {self.image_sum_service}")
            return
        
        if self.previous_prompt!='':
            prompt=self.previous_prompt
            logger.info(f"改变提示，使用上次提示")   

        if self.image_sum_service != "gemini":
            payload = {
                "model": model,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{base64_image}"
                                }
                            }
                        ]
                    }
                ],
                "max_tokens": 3000
            }
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}"
            }

        try:
            logger.info(f'handle_image: 请求图片内容{json.dumps(payload, ensure_ascii=False)}')
            response = requests.post(api_base, headers=headers, json=payload)
            logger.info(f"API请求已发送. 状态码: {response.status_code}")
            response.raise_for_status()
            logger.info("API响应状态码正常，开始解析JSON")
            response_json = response.json()
            logger.info("JSON解析完成")

            if self.image_sum_service == "gemini":
                reply_content = response_json.get('candidates', [{}])[0].get('content', {}).get('parts', [{}])[0].get('text', 'No text found in the response')
                logger.info(f"成功解析Gemini响应. 回复内容长度: {len(reply_content)}")
            else:
                if "choices" in response_json and len(response_json["choices"]) > 0:
                    first_choice = response_json["choices"][0]
                    if "message" in first_choice and "content" in first_choice["message"]:
                        response_content = first_choice["message"]["content"].strip()
                        logger.info("LLM API response content")
                        reply_content = response_content
                        if msg.ctype == ContextType.IMAGE and self.previous_prompt =='':
                            reply_content="您刚刚上传了一张图片，请问我有什么可以帮您的呢？"
                            self.previous_prompt=''
                        elif self.previous_prompt !='' and 'last_image_base64' in self.params_cache[user_id]:
                            self.previous_prompt =''
                            logger.info(f'last_image_base64 已经回答最后的提示,清空previous_prompt')
                    else:
                        logger.error("Content not found in the response")
                        reply_content = "Content not found in the LLM API response"
                else:
                    logger.error("No choices available in the response")
                    reply_content = "No choices available in the LLM API response"
        except Exception as e:
            logger.error(f"Error processing LLM API response: {e}")
            reply_content = f"An error occurred while processing LLM API response"

        reply = Reply()
        reply.type = ReplyType.TEXT
        # reply.content = f"{remove_markdown(reply_content)}\n\n💬5min内输入{self.image_sum_qa_prefix}+问题，可继续追问"
        reply.content = f"{remove_markdown(reply_content)}"
        e_context["reply"] = reply
        e_context.action = EventAction.BREAK_PASS
 

    def read_pdf(self, file_path):
        logger.info(f"开始读取PDF文件：{file_path}")
        doc = fitz.open(file_path)
        content = ' '.join([page.get_text() for page in doc])
        logger.info(f"PDF文件读取完成：{file_path}")

        return content
    def read_word(self, file_path):
        doc = Document(file_path)
        return ' '.join([p.text for p in doc.paragraphs])
    def read_markdown(self, file_path):
        with open(file_path, 'r', encoding='utf-8') as file:
            md_content = file.read()
            return markdown.markdown(md_content)
    def read_excel(self, file_path):
        workbook = load_workbook(file_path)
        content = ''
        for sheet in workbook:
            for row in sheet.iter_rows():
                content += ' '.join([str(cell.value) for cell in row])
                content += '\n'
        return content
    def read_txt(self, file_path):
        logger.debug(f"开始读取TXT文件: {file_path}")
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            logger.debug(f"TXT文件读取完成: {file_path}")
            logger.debug("TXT文件内容的前50个字符：")
            logger.debug(content[:50])  # 打印文件内容的前50个字符
            return content
        except Exception as e:
            logger.error(f"读取TXT文件时出错: {file_path}，错误信息: {str(e)}")
            return ""
    def read_csv(self, file_path):
        content = ''
        with open(file_path, 'r', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                content += ' '.join(row) + '\n'
        return content
    def read_html(self, file_path):
        with open(file_path, 'r', encoding='utf-8') as file:
            soup = BeautifulSoup(file, 'html.parser')
            return soup.get_text()
    def read_ppt(self, file_path):
        presentation = Presentation(file_path)
        content = ''
        for slide in presentation.slides:
            for shape in slide.shapes:
                if hasattr(shape, "text"):
                    content += shape.text + '\n'
        return content
    def extract_content(self, file_path):
        logger.info(f"extract_content: 提取文件内容，文件路径: {file_path}")
        file_size = os.path.getsize(file_path) // 1000  # 将文件大小转换为KB
        if file_size > int(self.max_file_size):
            logger.warning(f"文件大小超过限制({self.max_file_size}KB),不进行处理。文件大小: {file_size}KB")
            return None
        file_extension = os.path.splitext(file_path)[1][1:].lower()
        logger.info(f"extract_content: 文件类型为 {file_extension}")

        file_type = EXTENSION_TO_TYPE.get(file_extension)

        if not file_type:
            logger.error(f"不支持的文件扩展名: {file_extension}")
            return None

        read_func = {
            'pdf': self.read_pdf,
            'docx': self.read_word,
            'md': self.read_markdown,
            'txt': self.read_txt,
            'excel': self.read_excel,
            'csv': self.read_csv,
            'html': self.read_html,
            'ppt': self.read_ppt
        }.get(file_type)

        if not read_func:
            logger.error(f"不支持的文件类型: {file_type}")
            return None
        logger.info("extract_content: 文件内容提取完成")
        return read_func(file_path)
    def encode_image_to_base64(self, image_path):
        logger.info(f"开始处理图片: {image_path}")
        try:
            with Image.open(image_path) as img:
                logger.info(f"成功打开图片. 原始大小: {img.size}")
                if img.width > 1024:
                    new_size = (1024, int(img.height*1024/img.width))
                    img = img.resize(new_size)
                    img.save(image_path)  # 保存调整大小后的图片
                    logger.info(f"调整图片大小至: {new_size}")

            with open(image_path, "rb") as image_file:
                img_byte_arr = image_file.read()
                logger.info(f"读取图片完成. 大小: {len(img_byte_arr)} 字节")

            encoded = base64.b64encode(img_byte_arr).decode('ascii')
            logger.info(f"Base64编码完成. 编码后长度: {len(encoded)}")
            return encoded
        except Exception as e:
            logger.error(f"图片编码过程中发生错误: {str(e)}", exc_info=True)
            raise   
    
def remove_markdown(text):
        # 替换Markdown的粗体标记
        text = text.replace("**", "")
        # 替换Markdown的标题标记
        text = text.replace("### ", "").replace("## ", "").replace("# ", "")
        return text