ai_service.py

import logging
import json
import re
import mimetypes
import base64
from typing import Optional, List
import config

try:
    from openai import OpenAI

    HAS_OPENAI = True
except ImportError:
    HAS_OPENAI = False

logger = logging.getLogger("name_identify")


class AIService:
    """
    负责调用大语言模型进行文档/图片的名称识别。
    封装为类的好处：
    1. 可以维持配置状态（如 client1, client2 的持久化连接）。
    2. 可以方便地拓展不同的方法，如文本识别和图片识别，并共享回退和重试逻辑。
    """

    def __init__(self):
        self.client1 = None
        self.client2 = None
        if HAS_OPENAI:
            if config.AI_API_KEY_1 and config.AI_MODEL_1:
                self.client1 = OpenAI(
                    api_key=config.AI_API_KEY_1, base_url=config.AI_BASE_URL_1
                )
            if config.AI_API_KEY_2 and config.AI_MODEL_2:
                self.client2 = OpenAI(
                    api_key=config.AI_API_KEY_2, base_url=config.AI_BASE_URL_2
                )

    def get_name_from_content(
        self, text: str, is_book: bool, images_base64: List[str] = None
    ) -> Optional[str]:
        if not HAS_OPENAI:
            return None

        if not text.strip() and not images_base64:
            return None

        if not self.client1 and not self.client2:
            logger.warning("未配置任何 AI_API_KEY，跳过 AI 识别。")
            return None

        prompt = ""
        if is_book:
            prompt = (
                "以下是一本电子书（或扫描版文档/演示文稿）的前几页内容，请识别并返回这本书的书名或演示文稿的主题。\n"
                "如果能确定作者/演讲者，也请一并返回。\n"
                '请严格以 JSON 格式返回：{"title": "书名/主题", "author": "作者/演讲者"}\n'
                '如果无法确定，返回 json: {"title": ""}\n\n'
            )
        else:
            prompt = (
                "以下是一篇文档、演示文稿或音频文本的前几部分内容，请识别或总结出它的标题或主题。\n"
                '请严格以 JSON 格式返回：{"title": "标题/主题"}\n'
                '如果无法确定，返回 json: {"title": ""}\n\n'
            )

        # 构造用户输入内容
        user_content = []
        if text.strip():
            user_content.append(
                {"type": "text", "text": prompt + f"--- 文本内容 ---\n{text[:5000]}"}
            )
        else:
            user_content.append(
                {
                    "type": "text",
                    "text": prompt
                    + "--- 图像内容 --- (请从以下图片中读取封面标题或内页标题)",
                }
            )

        if images_base64:
            for img_b64 in images_base64:
                user_content.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"},
                    }
                )

        messages = [
            {
                "role": "system",
                "content": '你是一个高度准确的文档标题提取助手。你必须严格按照要求的 JSON 格式返回，不要包含任何其他说明文字，也不要使用 ```json 代码块。如果无法确定或找不到，必须严格返回 {"title": ""}',
            },
            {"role": "user", "content": user_content},
        ]

        return self._execute_completion(messages)

    def get_name_from_image(
        self, image_path: str, extra_prompt: str = ""
    ) -> Optional[str]:
        """专门针对图片的视觉识别"""
        if not HAS_OPENAI:
            return None

        if not self.client1 and not self.client2:
            logger.warning("未配置任何 AI_API_KEY，跳过 AI 识别。")
            return None

        try:
            with open(image_path, "rb") as f:
                base64_image = base64.b64encode(f.read()).decode("utf-8")
        except Exception as e:
            logger.error(f"  [图片读取失败] {e}")
            return None

        mime_type, _ = mimetypes.guess_type(image_path)
        if not mime_type:
            mime_type = "image/jpeg"

        prompt = (
            "请识别这张图片的主题或内容，并给出一个简短、描述性的文件名（不需要包含后缀名）。\n"
            '请严格以 JSON 格式返回：{"title": "图片文件名"}\n'
            '如果无法确定，返回 json: {"title": ""}\n'
        )
        if extra_prompt:
            prompt += f"\n额外提示：{extra_prompt}\n"

        messages = [
            {
                "role": "system",
                "content": '你是一个高度准确的图片命名助手。你必须严格按照要求的 JSON 格式返回，不要包含任何其他说明文字，也不要使用 ```json 代码块。如果无法确定或找不到，必须严格返回 {"title": ""}',
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:{mime_type};base64,{base64_image}"},
                    },
                ],
            },
        ]

        return self._execute_completion(messages)

    def _execute_completion(self, messages) -> Optional[str]:
        """执行大模型补全请求，包含主备容灾切换逻辑"""
        # 尝试客户端1
        if self.client1:
            try:
                response = self.client1.chat.completions.create(
                    model=config.AI_MODEL_1,
                    messages=messages,
                    temperature=0.3,
                )
                return self._parse_json_result(response.choices[0].message.content)
            except Exception as e:
                # 判断是否为严重且无法通过切换备用解决的本地配置错误
                # （如 API_KEY 格式明显错误等，但大多数网络/服务端/限流异常都值得尝试 fallback）
                logger.warning(f"  [主AI调用异常] {e}")

                if self.client2:
                    logger.info("  [AI备用切换] 尝试切换备用接口...")
                    return self._call_fallback(messages)
                else:
                    return None

        # 只有客户端2
        elif self.client2:
            return self._call_fallback(messages)

        return None

    def _call_fallback(self, messages) -> Optional[str]:
        try:
            response = self.client2.chat.completions.create(
                model=config.AI_MODEL_2, messages=messages, temperature=0.3
            )
            return self._parse_json_result(response.choices[0].message.content)
        except Exception as e:
            logger.error(f"  [备用AI调用异常] {e}")
            return None

    def _parse_json_result(self, text: str) -> Optional[str]:
        text = text.strip()

        # 1. 尝试提取 Markdown 代码块中的内容
        match = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
        json_text = match.group(1).strip() if match else text

        # 2. 如果存在杂乱文本，尝试截取首个 { 到最后一个 } 之间的内容
        if not json_text.startswith("{") and "{" in json_text and "}" in json_text:
            start_idx = json_text.find("{")
            end_idx = json_text.rfind("}") + 1
            json_text = json_text[start_idx:end_idx]

        try:
            data = json.loads(json_text)
            title = data.get("title", "").strip()
            author = data.get("author", "").strip()

            if not title:
                return None

            if author:
                return f"[{author}] {title}"
            return title
        except Exception as e:
            # === 非标准 JSON 输出的回退处理机制 ===
            logger.debug(f"JSON load failed, fallback. err: {e}")

            # 1. 明确表示无法识别的常见话术
            failure_keywords = [
                "抱歉",
                "没有找到",
                "无法确定",
                "无法识别",
                "未找到",
                "不知道",
                "未提供",
            ]
            if any(kw in text for kw in failure_keywords):
                logger.info("  [AI未找到内容] AI 回复表示无法识别标题。")
                return None

            # 2. 尝试使用正则提取书名号 《》 中的内容作为 fallback
            title_match = re.search(r"《(.*?)》", text)
            if title_match:
                extracted_title = title_match.group(1).strip()
                logger.info(f"  [正则回退提取成功] 从非 JSON 文本中提取到标题: {extracted_title}")
                return extracted_title

            logger.warning(f"  [JSON解析失败] 无法解析 AI 返回内容: {text[:100]}")
            return None

    def transcribe_audio(self, audio_file_path: str) -> Optional[str]:
        """使用 OpenAI Whisper 模型将音频转换为文本"""
        if not HAS_OPENAI:
            return None

        if not self.client1 and not self.client2:
            logger.warning("未配置任何 AI_API_KEY，跳过音频转录。")
            return None

        logger.info(f"  [AI音频转录] 正在转录音频: {audio_file_path}")

        # 尝试客户端1
        if self.client1:
            try:
                with open(audio_file_path, "rb") as audio_file:
                    # 默认使用 whisper-1 模型
                    response = self.client1.audio.transcriptions.create(
                        model="whisper-1",
                        file=audio_file,
                        response_format="text"
                    )
                return response.strip()
            except Exception as e:
                logger.warning(f"  [主AI音频转录异常] {e}")
                if self.client2:
                    logger.info("  [AI备用切换] 尝试使用备用接口转录音频...")
                    return self._transcribe_audio_fallback(audio_file_path)
                else:
                    return None

        # 只有客户端2
        elif self.client2:
            return self._transcribe_audio_fallback(audio_file_path)

        return None

    def _transcribe_audio_fallback(self, audio_file_path: str) -> Optional[str]:
        """备用音频转录"""
        try:
            with open(audio_file_path, "rb") as audio_file:
                response = self.client2.audio.transcriptions.create(
                    model="whisper-1",
                    file=audio_file,
                    response_format="text"
                )
            return response.strip()
        except Exception as e:
            logger.error(f"  [备用AI音频转录异常] {e}")
            return None


# 单例实例，供其他模块导入使用
ai_service = AIService()