modules/ai_model.py

# modules/ai_model.py
import torch
import base64
import requests
from io import BytesIO
import os
from huggingface_hub import login
from PIL import Image
from transformers import AutoProcessor, Gemma3nForConditionalGeneration
from utils.logger import log
from typing import Union, Tuple

class AIModel:
    def __init__(self, model_name: str = "google/gemma-3n-e2b-it"):
        self.model_name = model_name
        self.model = None
        self.processor = None
        
        # 设置缓存目录
        self._setup_cache_dirs()
        self._initialize_model()

    def _setup_cache_dirs(self):
        """设置缓存目录"""
        cache_dir = "/app/.cache/huggingface"
        os.makedirs(cache_dir, exist_ok=True)
        
        # 设置环境变量
        os.environ["HF_HOME"] = cache_dir
        os.environ["TRANSFORMERS_CACHE"] = cache_dir
        os.environ["HF_DATASETS_CACHE"] = cache_dir
        
        log.info(f"设置缓存目录: {cache_dir}")

    def _authenticate_hf(self):

        assitant_token = os.getenv("Assitant_tocken")
        token_to_use = assitant_token

        cache_dir = "/app/.cache/huggingface"
        login(token=token_to_use, add_to_git_credential=False)
        log.info("✅ HuggingFace 认证成功")
        return token_to_use



    def _initialize_model(self):
        """初始化Gemma模型"""
        try:
            log.info(f"正在加载模型: {self.model_name}")
            
            token = self._authenticate_hf()
            
            if not token:
                log.error("❌ 无法获取有效token，模型加载失败")
                self.model = None
                self.processor = None
                return
            
            cache_dir = "/app/.cache/huggingface"
            
            self.model = Gemma3nForConditionalGeneration.from_pretrained(
                self.model_name,
                device_map="auto",
                torch_dtype=torch.bfloat16,
                trust_remote_code=True,
                token=token,
                cache_dir=cache_dir
            ).eval()
            
            self.processor = AutoProcessor.from_pretrained(
                self.model_name,
                trust_remote_code=True,
                token=token,
                cache_dir=cache_dir
            )
            
            log.info("✅ Gemma AI 模型初始化成功")
            
        except Exception as e:
            log.error(f"❌ Gemma AI 模型初始化失败: {e}", exc_info=True)
            self.model = None
            self.processor = None

    def is_available(self) -> bool:
        
        return self.model is not None and self.processor is not None

    def detect_input_type(self, input_data: str) -> str:

        if not isinstance(input_data, str):
            return "text"
    
        image_extensions = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"]
        if (input_data.startswith(("http://", "https://")) and 
            any(input_data.lower().endswith(ext) for ext in image_extensions)):
            return "image"
        elif any(input_data.endswith(ext) for ext in image_extensions):
            return "image"
        elif input_data.startswith("data:image/"):
            return "image"

        audio_extensions = [".wav", ".mp3", ".m4a", ".ogg", ".flac"]
        if (input_data.startswith(("http://", "https://")) and 
            any(input_data.lower().endswith(ext) for ext in audio_extensions)):
            return "audio"
        elif any(input_data.endswith(ext) for ext in audio_extensions):
            return "audio"
    
        return "text"
    
    def transcribe_audio(self, audio_path: str) -> str:
        """
        使用 Hugging Face Inference API 将音频文件转写为文本。
        - 通过环境变量加载 HF_TOKEN 保证安全。
        - 包含网络请求超时和状态码检查，增强健壮性。
        """
        # 1. 从环境变量安全地获取 Token
        hf_token = os.getenv("Assitant_tocken")

        API_URL = "https://api-inference.huggingface.co/models/openai/whisper-large" # 建议使用更新的 v3 版本
        headers = {"Authorization": f"Bearer {hf_token}"}
    
        # 2. 检查音频文件是否存在
        if not os.path.exists(audio_path):
            log.error(f"❌ 音频文件不存在: {audio_path}")
            raise FileNotFoundError(f"指定的音频文件路径不存在: {audio_path}")

        try:
            with open(audio_path, "rb") as f:
                # 3. 发送请求，并设置较长的超时时间 (例如 60 秒)
                log.info(f"🎤 正在向 HF API 发送音频数据... (超时设置为60秒)")
                response = requests.post(API_URL, headers=headers, data=f, timeout=60)

            # 4. 检查 HTTP 响应状态码，主动抛出错误
            response.raise_for_status()  # 如果状态码不是 2xx，则会引发 HTTPError

            result = response.json()
            log.info("✅ HF API 响应成功。")

            # 5. 可靠地提取结果或处理错误信息
            if "text" in result:
                return result["text"].strip()
            else:
                error_message = result.get("error", "未知的 API 错误结构。")
                log.error(f"❌ 转录失败，API 返回: {error_message}")
                # 如果模型正在加载，HuggingFace 会在 error 字段中提示
                if isinstance(error_message, dict) and "estimated_time" in error_message:
                    raise RuntimeError(f"模型正在加载中，请稍后重试。预计等待时间: {error_message['estimated_time']:.1f}秒")
                raise RuntimeError(f"转录失败: {error_message}")

        except requests.exceptions.Timeout:
            log.error("❌ 请求超时！API 未在60秒内响应。")
            raise RuntimeError("语音识别服务请求超时，请稍后再试。")
        except requests.exceptions.RequestException as e:
            log.error(f"❌ 网络请求失败: {e}")
            raise RuntimeError(f"无法连接到语音识别服务: {e}")
        except Exception as e:
            # 捕获其他所有可能的异常，例如文件读取错误、JSON解码错误等
            log.error(f"❌ 处理音频时发生未知错误: {e}", exc_info=True)
            raise e


    def format_input(self, input_type: str, raw_input: str) -> Tuple[str, Union[str, Image.Image, None]]:
    
        if input_type == "image":
            try:
                if raw_input.startswith("data:image/"):
                    header, encoded = raw_input.split(",", 1)
                    image_data = base64.b64decode(encoded)
                    image = Image.open(BytesIO(image_data)).convert("RGB")
                elif raw_input.startswith(("http://", "https://")):
                    response = requests.get(raw_input, timeout=10)
                    response.raise_for_status()
                    image = Image.open(BytesIO(response.content)).convert("RGB")
                else:
                
                    image = Image.open(raw_input).convert("RGB")
            
                log.info("✅ 图片加载成功")
                return input_type, image, "请描述这张图片，并基于图片内容提供旅游建议。"
            
            except Exception as e:
                log.error(f"❌ 图片加载失败: {e}")
                return "text", None, f"图片加载失败，请检查路径或URL。"
            
        elif input_type == "audio":
            try:
                # --- 音频处理核心 ---
                # 假设: 您的类中有一个方法 `transcribe_audio` 用于语音转文字。
                # 您需要自行实现这个方法, 例如通过调用 Whisper, FunASR 或其他 ASR 服务。
                # 它接收音频文件路径 (raw_input) 并返回转写的文本字符串。
                log.info(f"🎤 开始处理音频文件: {raw_input}")
                transcribed_text = self.transcribe_audio(raw_input) 
                log.info(f"✅ 音频转写成功: '{transcribed_text[:50]}...'")
            
                # 注意：处理成功后，我们将 input_type 转为 "text"，
                # 因为音频内容已变为文本，后续流程可以统一处理。
                return "text", None, transcribed_text
            
            except Exception as e:
                log.error(f"❌ 音频处理失败: {e}", exc_info=True)
                return "text", None, f"音频处理失败，请检查文件或稍后再试。"
        
        else:  # text
            return input_type, None, raw_input

    def run_inference(self, input_type: str, formatted_input: Union[str, Image.Image], prompt: str,temperature: float = 0.5) -> str:
   
        try:
            inputs = self.processor(
                text=prompt,
                return_tensors="pt"
            ).to(self.model.device, dtype=torch.bfloat16)
            
            with torch.inference_mode():
                generation_args = {
                    "max_new_tokens": 1024,
                    "pad_token_id": self.processor.tokenizer.eos_token_id,
                    "use_cache": True
                }

                # 如果 temperature 接近0，使用贪心解码 (用于分类等确定性任务)
                if temperature < 1e-6: 
                    log.info("▶️ 使用贪心解码 (do_sample=False) 以获得确定性输出。")
                    generation_args["do_sample"] = False
                # 否则，使用采样解码 (用于创造性生成任务)
                else:
                    log.info(f"▶️ 使用采样解码 (do_sample=True)，temperature={temperature}。")
                    generation_args["do_sample"] = True
                    generation_args["temperature"] = temperature
                    generation_args["top_p"] = 0.9 # top_p 只在采样时有意义

                # 使用构建好的参数字典来调用 generate
                outputs = self.model.generate(
                    **inputs,
                    **generation_args
                )
            input_length = inputs.input_ids.shape[-1]
            generated_tokens = outputs[0][input_length:]
            decoded = self.processor.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()

            return decoded if decoded else "我理解了您的问题，请告诉我更多具体信息。"

        except RuntimeError as e:
            if "shape" in str(e):
                log.error(f"❌ Tensor形状错误: {e}")
                return "输入处理遇到问题，请尝试简化您的问题。"
            raise e
        except Exception as e:
            log.error(f"❌ 模型推理失败: {e}", exc_info=True)
            return "抱歉，处理您的请求时遇到技术问题。"
        
    def chat_completion(self, model: str, messages: list, **kwargs) -> str:
            
        if not self.is_available():
            log.error("模型未就绪，无法执行 chat_completion")
            if kwargs.get("response_format", {}).get("type") == "json_object":
                return '{"error": "Model not available"}'
            return "抱歉，AI 模型当前不可用。"

        full_prompt = "\n".join([msg.get("content", "") for msg in messages])

        temperature = kwargs.get("temperature", 0.6)

        if kwargs.get("response_format", {}).get("type") == "json_object":
            # 在 prompt 末尾添加指令，强制模型输出 JSON
            full_prompt += "\n\n请注意：你的回答必须是一个严格的、不含任何额外解释和代码块标记的 JSON 对象。"
            # 对于JSON生成任务，使用较低的 temperature 以获得更稳定、确定性的结构
            temperature = 0.1 

        log.debug(f"▶️ 执行 chat_completion (适配器), temperature={temperature}, prompt='{full_prompt[:100]}...'")

        return self.run_inference(
            input_type="text",
            formatted_input=None,
            prompt=full_prompt,
            temperature=temperature # 将处理后的 temperature 传递下去
        )


    def _build_prompt(self, processed_text: str, context: str = "") -> str:

        if context:
            return (
                f"你是一个专业的旅游助手。请基于以下背景信息，用中文友好地回答用户的问题。\n\n"
                f"--- 背景信息 ---\n{context}\n\n"
                f"--- 用户问题 ---\n{processed_text}\n\n"
                f"请提供专业、实用的旅游建议："
            )
        else:
            return (
                f"你是一个专业的旅游助手。请用中文友好地回答用户的问题。\n\n"
                f"用户问题：{processed_text}\n\n"
                f"请提供专业、实用的旅游建议："
            )

    def generate(self, user_input: str, context: str = "") -> str:
        """主要的生成方法 - 保持原有逻辑"""
        if not self.is_available():
            return "抱歉，AI 模型当前不可用，请稍后再试。"
    
        try:
            # 1. 检测输入类型
            input_type = self.detect_input_type(user_input)
            log.info(f"检测到输入类型: {input_type}")
        
            # 2. 格式化输入
            input_type, formatted_data, processed_text = self.format_input(input_type, user_input)
        
            # 3. 构建prompt - 使用你的原有结构
            prompt = self._build_prompt(processed_text, context)
        
            # 4. 执行推理
            return self.run_inference("text", formatted_data, prompt)
            
        except Exception as e:
            log.error(f"❌ 生成回复时发生错误: {e}", exc_info=True)
            return "抱歉，我在思考时遇到了点麻烦，请稍后再试。"