service.py

"""
使用 mistralai 官方 Python 库的服务类
"""
import os
from mistralai import Mistral
from mistralai.models import SystemMessage, UserMessage, AssistantMessage
import base64
from typing import List, Dict, Any, Union, Optional

class Service:
    """
    使用 mistralai 官方库的服务类，支持多模态内容
    """
    def __init__(self):
        """
        初始化服务类
        """
        self.model = "mistral-small-latest"  # 默认模型
        # 尝试从不同来源获取 API 密钥
        self.api_key = None
        
        # 1. 从环境变量获取
        self.api_key = os.environ.get("MISTRAL_API_KEY")
        
        # 2. 如果环境变量中没有，尝试从 Hugging Face hub secrets 获取
        if not self.api_key:
            try:
                from huggingface_hub import get_secret
                self.api_key = get_secret("MISTRAL_API_KEY")
            except Exception:
                pass
        self.headers = {
            "Authorization": "Bearer YOUR_API_KEY_HERE"  # 这将在使用前被替换
        }
        api_key = os.environ.get("MISTRAL_API_KEY", "")
        if not api_key:
            try:
                from huggingface_hub import get_secret
                api_key = get_secret("MISTRAL_API_KEY")
            except Exception:
                pass
        if not api_key:
            raise ValueError("API 密钥未设置。请设置 service.headers['Authorization'] = 'Bearer YOUR_API_KEY'")
            
        # 初始化客户端
        self.client = Mistral(api_key=api_key)
        
    def load_system_prompt(self, prompt_file: str) -> str:
        """
        加载系统提示文件
        
        Args:
            prompt_file: 系统提示文件路径
            
        Returns:
            文件内容
        """
        try:
            with open(prompt_file, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            print(f"加载系统提示文件失败: {e}")
            return ""
    
    def get_response(self, messages: List[Dict[str, Any]]) -> str:
        """
        从 Mistral API 获取响应
        
        Args:
            messages: 消息列表，包含角色和内容
            
        Returns:
            API 响应内容
        """
        try:
            # 发送请求
            response = self.client.chat.complete(
                model=self.model,
                messages=messages,  # 直接使用字典形式的消息
                stream=False  # 不使用流式响应
            )
            
            # 提取响应内容
            return response.choices[0].message.content
            
        except Exception as e:
            error_msg = f"API 请求错误: {str(e)}"
            print(error_msg)
            return error_msg
    
    def chat_with_image(self, 
                        text_prompt: str, 
                        image_base64: Optional[str] = None, 
                        history: Optional[List[Dict[str, Any]]] = None) -> str:
        """
        结合文本和图像（如果有）进行聊天
        
        Args:
            text_prompt: 文本提示
            image_base64: 图像的base64编码字符串（可选）
            history: 聊天历史记录（可选）
            
        Returns:
            模型响应
        """
        # 初始化消息列表
        if not history:
            history = []
            
        messages = list(history)  # 复制历史记录
        
        # 为当前用户请求创建消息内容
        user_content: Union[str, List[Dict[str, Any]]]
        
        if image_base64:
            # 如果有图像，创建多模态内容
            user_content = [
                {"type": "text", "text": text_prompt if text_prompt else "请分析这张图片"},
                {"type": "image_url", "image_url": {"url": image_base64}}
            ]
        else:
            # 纯文本内容
            user_content = text_prompt
            
        # 添加用户消息
        messages.append({"role": "user", "content": user_content})
        
        # 获取响应
        return self.get_response(messages)


# 示例用法
if __name__ == "__main__":
    # 创建服务实例
    service = Service()
    service.headers["Authorization"] = "Bearer YOUR_API_KEY"  # 替换为实际 API 密钥
    
    # 加载系统提示
    system_prompt = "你是一个有用的AI助手，可以回答问题和分析图像。"
    
    # 准备消息
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": "你好，请介绍一下自己"}
    ]
    
    # 获取响应
    response = service.get_response(messages)
    print(response)