ai-chat-system/services/ai_service.py

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

logger = logging.getLogger(__name__)

# 默认配置
DEFAULT_API_BASE = "http://192.168.2.17:19007/v1"
DEFAULT_API_KEY = "xxxx"
DEFAULT_MODEL = "auto"


class AIService:
    def __init__(self, api_base: str = None, api_key: str = None, model: str = None):
        self.api_base = api_base or DEFAULT_API_BASE
        self.api_key = api_key or DEFAULT_API_KEY
        self.model = model or DEFAULT_MODEL
    
    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
        logger.info(f"AI配置已更新: {api_base}, model={model}")
        self.api_base = api_base
        self.api_key = api_key
        self.model = model
    
    async def chat(self, messages: List[Dict], stream: bool = False) -> str:
        """
        调用AI模型进行对话
        
        Args:
            messages: 对话历史 [{"role": "user", "content": "..."}]
            stream: 是否流式输出
        
        Returns:
            AI回复内容
        """
        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": stream,
            "temperature": 0.7,
            "max_tokens": 2000
        }
        
        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']
    
    async def chat_stream(self, messages: List[Dict]):
        """
        流式调用AI模型
        """
        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()