mirror of https://github.com/Mai-with-u/MaiBot.git
421 lines
18 KiB
Python
421 lines
18 KiB
Python
import aiohttp
|
||
import asyncio
|
||
import requests
|
||
import time
|
||
import re
|
||
from typing import Tuple, Union
|
||
from nonebot import get_driver
|
||
from loguru import logger
|
||
from ..chat.config import global_config
|
||
from ..chat.utils_image import compress_base64_image_by_scale
|
||
|
||
driver = get_driver()
|
||
config = driver.config
|
||
|
||
|
||
class LLM_request:
|
||
def __init__(self, model, **kwargs):
|
||
# 将大写的配置键转换为小写并从config中获取实际值
|
||
try:
|
||
self.api_key = getattr(config, model["key"])
|
||
self.base_url = getattr(config, model["base_url"])
|
||
except AttributeError as e:
|
||
logger.error(f"配置错误：找不到对应的配置项 - {str(e)}")
|
||
raise ValueError(f"配置错误：找不到对应的配置项 - {str(e)}") from e
|
||
self.model_name = model["name"]
|
||
self.params = kwargs
|
||
|
||
async def generate_response(self, prompt: str) -> Tuple[str, str]:
|
||
"""根据输入的提示生成模型的异步响应"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
# 构建请求体
|
||
data = {
|
||
"model": self.model_name,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
**self.params
|
||
}
|
||
|
||
# 发送请求到完整的chat/completions端点
|
||
api_url = f"{self.base_url.rstrip('/')}/chat/completions"
|
||
logger.info(f"发送请求到URL: {api_url}+{self.model_name}") # 记录请求的URL
|
||
|
||
max_retries = 3
|
||
base_wait_time = 15
|
||
|
||
for retry in range(max_retries):
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(api_url, headers=headers, json=data) as response:
|
||
if response.status == 429:
|
||
wait_time = base_wait_time * (2 ** retry) # 指数退避
|
||
logger.warning(f"遇到请求限制(429)，等待{wait_time}秒后重试...")
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
|
||
if response.status in [500, 503]:
|
||
logger.error(f"服务器错误: {response.status}")
|
||
raise RuntimeError("服务器负载过高，模型恢复失败QAQ")
|
||
|
||
response.raise_for_status() # 检查其他响应状态
|
||
|
||
result = await response.json()
|
||
if "choices" in result and len(result["choices"]) > 0:
|
||
message = result["choices"][0]["message"]
|
||
content = message.get("content", "")
|
||
think_match = None
|
||
reasoning_content = message.get("reasoning_content", "")
|
||
if not reasoning_content:
|
||
think_match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
|
||
if think_match:
|
||
reasoning_content = think_match.group(1).strip()
|
||
content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
|
||
return content, reasoning_content
|
||
return "没有返回结果", ""
|
||
|
||
except Exception as e:
|
||
if retry < max_retries - 1: # 如果还有重试机会
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.error(f"[回复]请求失败，等待{wait_time}秒后重试... 错误: {str(e)}", exc_info=True)
|
||
await asyncio.sleep(wait_time)
|
||
else:
|
||
logger.critical(f"请求失败: {str(e)}", exc_info=True)
|
||
raise RuntimeError(f"API请求失败: {str(e)}")
|
||
|
||
logger.error("达到最大重试次数，请求仍然失败")
|
||
raise RuntimeError("达到最大重试次数，API请求仍然失败")
|
||
|
||
async def generate_response_for_image(self, prompt: str, image_base64: str) -> Tuple[str, str]:
|
||
"""根据输入的提示和图片生成模型的异步响应"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
# 构建请求体
|
||
def build_request_data(img_base64: str):
|
||
return {
|
||
"model": self.model_name,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": prompt
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{img_base64}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
],
|
||
**self.params
|
||
}
|
||
|
||
|
||
# 发送请求到完整的chat/completions端点
|
||
api_url = f"{self.base_url.rstrip('/')}/chat/completions"
|
||
logger.info(f"发送请求到URL: {api_url}+{self.model_name}") # 记录请求的URL
|
||
|
||
max_retries = 3
|
||
base_wait_time = 15
|
||
|
||
current_image_base64 = image_base64
|
||
current_image_base64 = compress_base64_image_by_scale(current_image_base64)
|
||
|
||
|
||
for retry in range(max_retries):
|
||
try:
|
||
data = build_request_data(current_image_base64)
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(api_url, headers=headers, json=data) as response:
|
||
if response.status == 429:
|
||
wait_time = base_wait_time * (2 ** retry) # 指数退避
|
||
logger.warning(f"遇到请求限制(429)，等待{wait_time}秒后重试...")
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
|
||
elif response.status == 413:
|
||
logger.warning("图片太大(413)，尝试压缩...")
|
||
current_image_base64 = compress_base64_image_by_scale(current_image_base64)
|
||
continue
|
||
|
||
response.raise_for_status() # 检查其他响应状态
|
||
|
||
result = await response.json()
|
||
if "choices" in result and len(result["choices"]) > 0:
|
||
message = result["choices"][0]["message"]
|
||
content = message.get("content", "")
|
||
think_match = None
|
||
reasoning_content = message.get("reasoning_content", "")
|
||
if not reasoning_content:
|
||
think_match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
|
||
if think_match:
|
||
reasoning_content = think_match.group(1).strip()
|
||
content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
|
||
return content, reasoning_content
|
||
return "没有返回结果", ""
|
||
|
||
except Exception as e:
|
||
if retry < max_retries - 1: # 如果还有重试机会
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.error(f"[image回复]请求失败，等待{wait_time}秒后重试... 错误: {str(e)}", exc_info=True)
|
||
await asyncio.sleep(wait_time)
|
||
else:
|
||
logger.critical(f"请求失败: {str(e)}", exc_info=True)
|
||
raise RuntimeError(f"API请求失败: {str(e)}")
|
||
|
||
logger.error("达到最大重试次数，请求仍然失败")
|
||
raise RuntimeError("达到最大重试次数，API请求仍然失败")
|
||
|
||
async def generate_response_async(self, prompt: str) -> Union[str, Tuple[str, str]]:
|
||
"""异步方式根据输入的提示生成模型的响应"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
# 构建请求体
|
||
data = {
|
||
"model": self.model_name,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"temperature": 0.5,
|
||
**self.params
|
||
}
|
||
|
||
# 发送请求到完整的 chat/completions 端点
|
||
api_url = f"{self.base_url.rstrip('/')}/chat/completions"
|
||
logger.info(f"Request URL: {api_url}") # 记录请求的 URL
|
||
|
||
max_retries = 3
|
||
base_wait_time = 15
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
for retry in range(max_retries):
|
||
try:
|
||
async with session.post(api_url, headers=headers, json=data) as response:
|
||
if response.status == 429:
|
||
wait_time = base_wait_time * (2 ** retry) # 指数退避
|
||
logger.warning(f"遇到请求限制(429)，等待{wait_time}秒后重试...")
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
|
||
response.raise_for_status() # 检查其他响应状态
|
||
|
||
result = await response.json()
|
||
if "choices" in result and len(result["choices"]) > 0:
|
||
content = result["choices"][0]["message"]["content"]
|
||
reasoning_content = result["choices"][0]["message"].get("reasoning_content", "")
|
||
return content, reasoning_content
|
||
return "没有返回结果", ""
|
||
|
||
except Exception as e:
|
||
if retry < max_retries - 1: # 如果还有重试机会
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.error(f"[回复]请求失败，等待{wait_time}秒后重试... 错误: {str(e)}")
|
||
await asyncio.sleep(wait_time)
|
||
else:
|
||
logger.error(f"请求失败: {str(e)}")
|
||
return f"请求失败: {str(e)}", ""
|
||
|
||
logger.error("达到最大重试次数，请求仍然失败")
|
||
return "达到最大重试次数，请求仍然失败", ""
|
||
|
||
|
||
|
||
def generate_response_for_image_sync(self, prompt: str, image_base64: str) -> Tuple[str, str]:
|
||
"""同步方法：根据输入的提示和图片生成模型的响应"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
image_base64=compress_base64_image_by_scale(image_base64)
|
||
|
||
# 构建请求体
|
||
data = {
|
||
"model": self.model_name,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": prompt
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{image_base64}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
],
|
||
**self.params
|
||
}
|
||
|
||
# 发送请求到完整的chat/completions端点
|
||
api_url = f"{self.base_url.rstrip('/')}/chat/completions"
|
||
logger.info(f"发送请求到URL: {api_url}+{self.model_name}") # 记录请求的URL
|
||
|
||
max_retries = 2
|
||
base_wait_time = 6
|
||
|
||
for retry in range(max_retries):
|
||
try:
|
||
response = requests.post(api_url, headers=headers, json=data, timeout=30)
|
||
|
||
if response.status_code == 429:
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.warning(f"遇到请求限制(429)，等待{wait_time}秒后重试...")
|
||
time.sleep(wait_time)
|
||
continue
|
||
|
||
response.raise_for_status() # 检查其他响应状态
|
||
|
||
result = response.json()
|
||
if "choices" in result and len(result["choices"]) > 0:
|
||
message = result["choices"][0]["message"]
|
||
content = message.get("content", "")
|
||
think_match = None
|
||
reasoning_content = message.get("reasoning_content", "")
|
||
if not reasoning_content:
|
||
think_match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
|
||
if think_match:
|
||
reasoning_content = think_match.group(1).strip()
|
||
content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
|
||
return content, reasoning_content
|
||
return "没有返回结果", ""
|
||
|
||
except Exception as e:
|
||
if retry < max_retries - 1: # 如果还有重试机会
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.error(f"[image_sync回复]请求失败，等待{wait_time}秒后重试... 错误: {str(e)}", exc_info=True)
|
||
time.sleep(wait_time)
|
||
else:
|
||
logger.critical(f"请求失败: {str(e)}", exc_info=True)
|
||
raise RuntimeError(f"API请求失败: {str(e)}")
|
||
|
||
logger.error("达到最大重试次数，请求仍然失败")
|
||
raise RuntimeError("达到最大重试次数，API请求仍然失败")
|
||
|
||
def get_embedding_sync(self, text: str, model: str = "BAAI/bge-m3") -> Union[list, None]:
|
||
"""同步方法：获取文本的embedding向量
|
||
|
||
Args:
|
||
text: 需要获取embedding的文本
|
||
model: 使用的模型名称，默认为"BAAI/bge-m3"
|
||
|
||
Returns:
|
||
list: embedding向量，如果失败则返回None
|
||
"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
data = {
|
||
"model": model,
|
||
"input": text,
|
||
"encoding_format": "float"
|
||
}
|
||
|
||
api_url = f"{self.base_url.rstrip('/')}/embeddings"
|
||
logger.info(f"发送请求到URL: {api_url}+{self.model_name}") # 记录请求的URL
|
||
|
||
max_retries = 2
|
||
base_wait_time = 6
|
||
|
||
for retry in range(max_retries):
|
||
try:
|
||
response = requests.post(api_url, headers=headers, json=data, timeout=30)
|
||
|
||
if response.status_code == 429:
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.warning(f"遇到请求限制(429)，等待{wait_time}秒后重试...")
|
||
time.sleep(wait_time)
|
||
continue
|
||
|
||
response.raise_for_status()
|
||
|
||
result = response.json()
|
||
if 'data' in result and len(result['data']) > 0:
|
||
return result['data'][0]['embedding']
|
||
return None
|
||
|
||
except Exception as e:
|
||
if retry < max_retries - 1:
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.error(f"[embedding_sync]请求失败，等待{wait_time}秒后重试... 错误: {str(e)}", exc_info=True)
|
||
time.sleep(wait_time)
|
||
else:
|
||
logger.critical(f"embedding请求失败: {str(e)}", exc_info=True)
|
||
return None
|
||
|
||
logger.error("达到最大重试次数，embedding请求仍然失败")
|
||
return None
|
||
|
||
async def get_embedding(self, text: str, model: str = "BAAI/bge-m3") -> Union[list, None]:
|
||
"""异步方法：获取文本的embedding向量
|
||
|
||
Args:
|
||
text: 需要获取embedding的文本
|
||
model: 使用的模型名称，默认为"BAAI/bge-m3"
|
||
|
||
Returns:
|
||
list: embedding向量，如果失败则返回None
|
||
"""
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
data = {
|
||
"model": model,
|
||
"input": text,
|
||
"encoding_format": "float"
|
||
}
|
||
|
||
api_url = f"{self.base_url.rstrip('/')}/embeddings"
|
||
logger.info(f"发送请求到URL: {api_url}+{self.model_name}") # 记录请求的URL
|
||
|
||
max_retries = 3
|
||
base_wait_time = 15
|
||
|
||
for retry in range(max_retries):
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(api_url, headers=headers, json=data) as response:
|
||
if response.status == 429:
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.warning(f"遇到请求限制(429)，等待{wait_time}秒后重试...")
|
||
await asyncio.sleep(wait_time)
|
||
continue
|
||
|
||
response.raise_for_status()
|
||
|
||
result = await response.json()
|
||
if 'data' in result and len(result['data']) > 0:
|
||
return result['data'][0]['embedding']
|
||
return None
|
||
|
||
except Exception as e:
|
||
if retry < max_retries - 1:
|
||
wait_time = base_wait_time * (2 ** retry)
|
||
logger.error(f"[embedding]请求失败，等待{wait_time}秒后重试... 错误: {str(e)}", exc_info=True)
|
||
await asyncio.sleep(wait_time)
|
||
else:
|
||
logger.critical(f"embedding请求失败: {str(e)}", exc_info=True)
|
||
return None
|
||
|
||
logger.error("达到最大重试次数，embedding请求仍然失败")
|
||
return None
|