ai-chat-system/services/ai_service.py

"""
AI服务 - 调用大模型API
"""
import httpx
from typing import List, Dict, AsyncGenerator
import json
import logging

logger = logging.getLogger(__name__)


class AIService:
    def __init__(self):
        self.api_base = ""
        self.api_key = ""
        self.model = ""
        self.use_mock = True
    
    def update_config(self, api_base: str, api_key: str, model: str):
        """更新配置"""
        self.api_base = api_base
        self.api_key = api_key
        self.model = model
        # 如果配置完整则使用真实API，否则使用mock
        self.use_mock = not (api_base and model)
        logger.info(f"AI配置已更新: api_base={api_base}, model={model}, use_mock={self.use_mock}")
    
    async def chat(self, messages: List[Dict]) -> str:
        """
        调用AI模型进行对话
        """
        # 如果使用mock模式，返回模拟回复
        if self.use_mock:
            logger.info("使用Mock模式回复")
            last_msg = messages[-1]['content'] if messages else "你好"
            return f"这是一个测试回复。您说的是：{last_msg}\n\n请配置有效的AI服务地址和模型，才能获得真正的AI回复。"
        
        # 调用真实API
        url = f"{self.api_base}/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 2000
        }
        
        logger.info(f"调用AI API: {url}, model={self.model}")
        
        try:
            async with httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(url, headers=headers, json=payload)
                response.raise_for_status()
                data = response.json()
                return data['choices'][0]['message']['content']
        except Exception as e:
            logger.error(f"AI API调用失败: {e}")
            # API失败时返回模拟回复
            last_msg = messages[-1]['content'] if messages else "你好"
            return f"AI服务暂时不可用（错误：{str(e)}）。您说的是：{last_msg}"
    
    async def chat_stream(self, messages: List[Dict]) -> AsyncGenerator[str, None]:
        """
        流式调用AI模型
        """
        if self.use_mock:
            last_msg = messages[-1]['content'] if messages else "你好"
            reply = f"这是一个测试回复。您说的是：{last_msg}"
            for char in reply:
                yield char
            return
        
        url = f"{self.api_base}/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model,
            "messages": messages,
            "stream": True,
            "temperature": 0.7,
            "max_tokens": 2000
        }
        
        async with httpx.AsyncClient(timeout=60.0) as client:
            async with client.stream("POST", url, headers=headers, json=payload) as response:
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        data_str = line[6:]
                        if data_str == "[DONE]":
                            break
                        try:
                            data = json.loads(data_str)
                            if 'choices' in data and len(data['choices']) > 0:
                                delta = data['choices'][0].get('delta', {})
                                if 'content' in delta:
                                    yield delta['content']
                        except json.JSONDecodeError:
                            continue


# 全局实例
ai_service = AIService()