f395309207
- Add providers.py with 5 provider presets (OpenAI, DeepSeek, Xiaomi MiMo, Alibaba DashScope, MiniMax) - Add LLM_PROVIDER env var for one-line provider switching - Improve <think> tag stripping for reasoning models - Add .env.example with documented configuration - Update README with provider configuration section
142 lines
4.4 KiB
Python
142 lines
4.4 KiB
Python
"""
|
||
LLM客户端封装
|
||
统一使用OpenAI格式调用，支持提供商预设（DeepSeek、Xiaomi MiMo等）
|
||
"""
|
||
|
||
import json
|
||
import re
|
||
from typing import Optional, Dict, Any, List
|
||
from openai import OpenAI
|
||
|
||
from ..config import Config
|
||
|
||
|
||
# <think>标签的正则表达式
|
||
# 匹配 <think>...</think> 标签及其内容（支持多行，非贪婪匹配）
|
||
# 也处理 <think>...</think> 变体（某些模型可能使用略微不同的格式）
|
||
THINK_TAG_PATTERN = re.compile(r'<think>[\s\S]*?</think>\s*', re.IGNORECASE)
|
||
|
||
# 某些提供商的推理模型会在响应中包含<think>标签
|
||
PROVIDERS_WITH_THINK_TAGS = {"deepseek", "xiaomi_mimo", "minimax"}
|
||
|
||
|
||
class LLMClient:
|
||
"""LLM客户端，支持提供商预设配置"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: Optional[str] = None,
|
||
base_url: Optional[str] = None,
|
||
model: Optional[str] = None,
|
||
provider: Optional[str] = None,
|
||
):
|
||
self.api_key = api_key or Config.LLM_API_KEY
|
||
self.base_url = base_url or Config.LLM_BASE_URL
|
||
self.model = model or Config.LLM_MODEL_NAME
|
||
self.provider = provider or Config.LLM_PROVIDER
|
||
|
||
if not self.api_key:
|
||
raise ValueError("LLM_API_KEY 未配置")
|
||
|
||
self.client = OpenAI(
|
||
api_key=self.api_key,
|
||
base_url=self.base_url
|
||
)
|
||
|
||
# 判断是否需要强制清理<think>标签
|
||
# 如果是已知的推理模型提供商，总是清理；否则也清理（安全兜底）
|
||
self._should_strip_think = (
|
||
self.provider in PROVIDERS_WITH_THINK_TAGS
|
||
or True # 总是清理，因为不影响正常输出
|
||
)
|
||
|
||
def _strip_think_tags(self, content: str) -> str:
|
||
"""
|
||
移除响应中的<think>思考内容标签。
|
||
|
||
某些模型（如DeepSeek Reasoner、Xiaomi MiMo、MiniMax M2.5）
|
||
会在响应中包含<think>...</think>标签，需要移除以获得纯净输出。
|
||
"""
|
||
if not content:
|
||
return content
|
||
return THINK_TAG_PATTERN.sub('', content).strip()
|
||
|
||
def chat(
|
||
self,
|
||
messages: List[Dict[str, str]],
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 4096,
|
||
response_format: Optional[Dict] = None
|
||
) -> str:
|
||
"""
|
||
发送聊天请求
|
||
|
||
Args:
|
||
messages: 消息列表
|
||
temperature: 温度参数
|
||
max_tokens: 最大token数
|
||
response_format: 响应格式（如JSON模式）
|
||
|
||
Returns:
|
||
模型响应文本
|
||
"""
|
||
kwargs = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
}
|
||
|
||
if response_format:
|
||
kwargs["response_format"] = response_format
|
||
|
||
response = self.client.chat.completions.create(**kwargs)
|
||
content = response.choices[0].message.content or ""
|
||
|
||
# 移除<think>思考内容标签
|
||
content = self._strip_think_tags(content)
|
||
|
||
return content
|
||
|
||
def chat_json(
|
||
self,
|
||
messages: List[Dict[str, str]],
|
||
temperature: float = 0.3,
|
||
max_tokens: int = 4096
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
发送聊天请求并返回JSON
|
||
|
||
Args:
|
||
messages: 消息列表
|
||
temperature: 温度参数
|
||
max_tokens: 最大token数
|
||
|
||
Returns:
|
||
解析后的JSON对象
|
||
"""
|
||
response = self.chat(
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
response_format={"type": "json_object"}
|
||
)
|
||
# 清理markdown代码块标记
|
||
cleaned_response = response.strip()
|
||
cleaned_response = re.sub(r'^```(?:json)?\s*\n?', '', cleaned_response, flags=re.IGNORECASE)
|
||
cleaned_response = re.sub(r'\n?```\s*$', '', cleaned_response)
|
||
cleaned_response = cleaned_response.strip()
|
||
|
||
try:
|
||
return json.loads(cleaned_response)
|
||
except json.JSONDecodeError:
|
||
raise ValueError(f"LLM返回的JSON格式无效: {cleaned_response}")
|
||
|
||
def get_info(self) -> Dict[str, Any]:
|
||
"""返回当前客户端配置信息（用于日志/调试）"""
|
||
return {
|
||
"provider": self.provider or "custom",
|
||
"base_url": self.base_url,
|
||
"model": self.model,
|
||
}
|