mirror of
https://github.com/LiteyukiStudio/nonebot-plugin-marshoai.git
synced 2024-11-30 09:14:52 +08:00
323 lines
13 KiB
Python
323 lines
13 KiB
Python
import contextlib
|
||
import traceback
|
||
from typing import Optional
|
||
from pathlib import Path
|
||
|
||
from arclet.alconna import Alconna, Args, AllParam
|
||
from azure.ai.inference.models import (
|
||
UserMessage,
|
||
AssistantMessage,
|
||
ToolMessage,
|
||
TextContentItem,
|
||
ImageContentItem,
|
||
ImageUrl,
|
||
CompletionsFinishReason,
|
||
ChatCompletionsToolCall,
|
||
)
|
||
from azure.core.credentials import AzureKeyCredential
|
||
from nonebot import on_command, on_message, logger
|
||
from nonebot.adapters import Message, Event
|
||
from nonebot.params import CommandArg
|
||
from nonebot.permission import SUPERUSER
|
||
from nonebot.rule import Rule, to_me
|
||
from nonebot_plugin_alconna import on_alconna, MsgTarget
|
||
from nonebot_plugin_alconna.uniseg import UniMessage, UniMsg
|
||
import nonebot_plugin_localstore as store
|
||
from nonebot import get_driver
|
||
|
||
from .constants import *
|
||
from .metadata import metadata
|
||
from .models import MarshoContext, MarshoTools
|
||
from .util import *
|
||
|
||
|
||
async def at_enable():
|
||
return config.marshoai_at
|
||
|
||
|
||
driver = get_driver()
|
||
|
||
changemodel_cmd = on_command("changemodel", permission=SUPERUSER, priority=10, block=True)
|
||
resetmem_cmd = on_command("reset", priority=10, block=True)
|
||
# setprompt_cmd = on_command("prompt",permission=SUPERUSER)
|
||
praises_cmd = on_command("praises", permission=SUPERUSER, priority=10, block=True)
|
||
add_usermsg_cmd = on_command("usermsg", permission=SUPERUSER, priority=10, block=True)
|
||
add_assistantmsg_cmd = on_command("assistantmsg", permission=SUPERUSER, priority=10, block=True)
|
||
contexts_cmd = on_command("contexts", permission=SUPERUSER, priority=10, block=True)
|
||
save_context_cmd = on_command("savecontext", permission=SUPERUSER, priority=10, block=True)
|
||
load_context_cmd = on_command("loadcontext", permission=SUPERUSER, priority=10, block=True)
|
||
marsho_cmd = on_alconna(
|
||
Alconna(
|
||
config.marshoai_default_name,
|
||
Args["text?", AllParam],
|
||
),
|
||
aliases=config.marshoai_aliases,
|
||
priority=10,
|
||
block=True
|
||
)
|
||
marsho_at = on_message(rule=to_me()&at_enable, priority=11)
|
||
nickname_cmd = on_alconna(
|
||
Alconna(
|
||
"nickname",
|
||
Args["name?", str],
|
||
),
|
||
priority = 10,
|
||
block = True
|
||
)
|
||
refresh_data_cmd = on_command("refresh_data", permission=SUPERUSER, priority=10, block=True)
|
||
|
||
command_start = driver.config.command_start
|
||
model_name = config.marshoai_default_model
|
||
context = MarshoContext()
|
||
tools = MarshoTools()
|
||
token = config.marshoai_token
|
||
endpoint = config.marshoai_azure_endpoint
|
||
client = ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(token))
|
||
target_list = [] # 记录需保存历史上下文的列表
|
||
|
||
|
||
@driver.on_startup
|
||
async def _preload_tools():
|
||
tools_dir = store.get_plugin_data_dir() / "tools"
|
||
os.makedirs(tools_dir, exist_ok=True)
|
||
if config.marshoai_load_builtin_tools:
|
||
tools.load_tools(Path(__file__).parent / "tools")
|
||
tools.load_tools(store.get_plugin_data_dir() / "tools")
|
||
|
||
|
||
@add_usermsg_cmd.handle()
|
||
async def add_usermsg(target: MsgTarget, arg: Message = CommandArg()):
|
||
if msg := arg.extract_plain_text():
|
||
context.append(UserMessage(content=msg).as_dict(), target.id, target.private)
|
||
await add_usermsg_cmd.finish("已添加用户消息")
|
||
|
||
|
||
@add_assistantmsg_cmd.handle()
|
||
async def add_assistantmsg(target: MsgTarget, arg: Message = CommandArg()):
|
||
if msg := arg.extract_plain_text():
|
||
context.append(
|
||
AssistantMessage(content=msg).as_dict(), target.id, target.private
|
||
)
|
||
await add_assistantmsg_cmd.finish("已添加助手消息")
|
||
|
||
|
||
@praises_cmd.handle()
|
||
async def praises():
|
||
# await UniMessage(await tools.call("marshoai-weather.get_weather", {"location":"杭州"})).send()
|
||
await praises_cmd.finish(build_praises())
|
||
|
||
|
||
@contexts_cmd.handle()
|
||
async def contexts(target: MsgTarget):
|
||
backup_context = await get_backup_context(target.id, target.private)
|
||
if backup_context:
|
||
context.set_context(backup_context, target.id, target.private) # 加载历史记录
|
||
await contexts_cmd.finish(str(context.build(target.id, target.private)))
|
||
|
||
|
||
@save_context_cmd.handle()
|
||
async def save_context(target: MsgTarget, arg: Message = CommandArg()):
|
||
contexts_data = context.build(target.id, target.private)
|
||
if not context:
|
||
await save_context_cmd.finish("暂无上下文可以保存")
|
||
if msg := arg.extract_plain_text():
|
||
await save_context_to_json(msg, contexts_data, "contexts")
|
||
await save_context_cmd.finish("已保存上下文")
|
||
|
||
|
||
@load_context_cmd.handle()
|
||
async def load_context(target: MsgTarget, arg: Message = CommandArg()):
|
||
if msg := arg.extract_plain_text():
|
||
await get_backup_context(target.id, target.private) # 为了将当前会话添加到"已恢复过备份"的列表而添加，防止上下文被覆盖（好奇怪QwQ
|
||
context.set_context(
|
||
await load_context_from_json(msg, "contexts"), target.id, target.private
|
||
)
|
||
await load_context_cmd.finish("已加载并覆盖上下文")
|
||
|
||
|
||
@resetmem_cmd.handle()
|
||
async def resetmem(target: MsgTarget):
|
||
if [target.id, target.private] not in target_list:
|
||
target_list.append([target.id, target.private])
|
||
context.reset(target.id, target.private)
|
||
await resetmem_cmd.finish("上下文已重置")
|
||
|
||
|
||
@changemodel_cmd.handle()
|
||
async def changemodel(arg: Message = CommandArg()):
|
||
global model_name
|
||
if model := arg.extract_plain_text():
|
||
model_name = model
|
||
await changemodel_cmd.finish("已切换")
|
||
|
||
|
||
@nickname_cmd.handle()
|
||
async def nickname(event: Event, name=None):
|
||
nicknames = await get_nicknames()
|
||
user_id = event.get_user_id()
|
||
if not name:
|
||
if user_id not in nicknames:
|
||
await nickname_cmd.finish("你未设置昵称")
|
||
await nickname_cmd.finish("你的昵称为：" + str(nicknames[user_id]))
|
||
if name == "reset":
|
||
await set_nickname(user_id, "")
|
||
await nickname_cmd.finish("已重置昵称")
|
||
else:
|
||
await set_nickname(user_id, name)
|
||
await nickname_cmd.finish("已设置昵称为：" + name)
|
||
|
||
|
||
@refresh_data_cmd.handle()
|
||
async def refresh_data():
|
||
await refresh_nickname_json()
|
||
await refresh_praises_json()
|
||
await refresh_data_cmd.finish("已刷新数据")
|
||
|
||
|
||
@marsho_at.handle()
|
||
@marsho_cmd.handle()
|
||
async def marsho(target: MsgTarget, event: Event, text: Optional[UniMsg] = None):
|
||
global target_list
|
||
if event.get_message().extract_plain_text() and (not text and event.get_message().extract_plain_text() != config.marshoai_default_name):
|
||
text = event.get_message()
|
||
if not text:
|
||
# 发送说明
|
||
await UniMessage(metadata.usage + "\n当前使用的模型：" + model_name).send()
|
||
await marsho_cmd.finish(INTRODUCTION)
|
||
try:
|
||
user_id = event.get_user_id()
|
||
nicknames = await get_nicknames()
|
||
user_nickname = nicknames.get(user_id, "")
|
||
if user_nickname != "":
|
||
nickname_prompt = f"\n*此消息的说话者:{user_nickname}*"
|
||
else:
|
||
nickname_prompt = ""
|
||
# 用户名无法获取，暂时注释
|
||
# user_nickname = event.sender.nickname # 未设置昵称时获取用户名
|
||
# nickname_prompt = f"\n*此消息的说话者:{user_nickname}"
|
||
if config.marshoai_enable_nickname_tip:
|
||
await UniMessage(
|
||
"*你未设置自己的昵称。推荐使用'nickname [昵称]'命令设置昵称来获得个性化(可能）回答。"
|
||
).send()
|
||
|
||
is_support_image_model = model_name.lower() in SUPPORT_IMAGE_MODELS + config.marshoai_additional_image_models
|
||
is_reasoning_model = model_name.lower() in REASONING_MODELS
|
||
usermsg = [] if is_support_image_model else ""
|
||
for i in text:
|
||
if i.type == "text":
|
||
if is_support_image_model:
|
||
usermsg += [TextContentItem(text=i.data["text"] + nickname_prompt)]
|
||
else:
|
||
usermsg += str(i.data["text"] + nickname_prompt)
|
||
elif i.type == "image":
|
||
if is_support_image_model:
|
||
usermsg.append(
|
||
ImageContentItem(
|
||
image_url=ImageUrl(url=str(await get_image_b64(i.data["url"])))
|
||
)
|
||
)
|
||
elif config.marshoai_enable_support_image_tip:
|
||
await UniMessage("*此模型不支持图片处理。").send()
|
||
backup_context = await get_backup_context(target.id, target.private)
|
||
if backup_context:
|
||
context.set_context(backup_context, target.id, target.private) # 加载历史记录
|
||
logger.info(f"已恢复会话 {target.id} 的上下文备份~")
|
||
context_msg = context.build(target.id, target.private)
|
||
if not is_reasoning_model:
|
||
context_msg = [get_prompt()] + context_msg
|
||
# o1等推理模型不支持系统提示词, 故不添加
|
||
response = await make_chat(
|
||
client=client,
|
||
model_name=model_name,
|
||
msg=context_msg + [UserMessage(content=usermsg)],
|
||
tools=tools.get_tools_list()
|
||
)
|
||
# await UniMessage(str(response)).send()
|
||
choice = response.choices[0]
|
||
if (choice["finish_reason"] == CompletionsFinishReason.STOPPED): # 当对话成功时，将dict的上下文添加到上下文类中
|
||
context.append(
|
||
UserMessage(content=usermsg).as_dict(), target.id, target.private
|
||
)
|
||
context.append(choice.message.as_dict(), target.id, target.private)
|
||
if [target.id, target.private] not in target_list:
|
||
target_list.append([target.id, target.private])
|
||
await UniMessage(str(choice.message.content)).send(reply_to=True)
|
||
elif choice["finish_reason"] == CompletionsFinishReason.CONTENT_FILTERED:
|
||
await UniMessage("*已被内容过滤器过滤。请调整聊天内容后重试。").send(reply_to=True)
|
||
return
|
||
elif choice["finish_reason"] == CompletionsFinishReason.TOOL_CALLS:
|
||
tool_msg = []
|
||
while choice.message.tool_calls != None:
|
||
tool_msg.append(AssistantMessage(tool_calls=response.choices[0].message.tool_calls))
|
||
for tool_call in choice.message.tool_calls:
|
||
if isinstance(tool_call, ChatCompletionsToolCall):
|
||
function_args = json.loads(tool_call.function.arguments.replace("'", '"'))
|
||
logger.info(f"调用函数 {tool_call.function.name} ,参数为 {function_args}")
|
||
await UniMessage(f"调用函数 {tool_call.function.name} ,参数为 {function_args}").send()
|
||
func_return = await tools.call(tool_call.function.name, function_args)
|
||
tool_msg.append(ToolMessage(tool_call_id=tool_call.id, content=func_return))
|
||
response = await make_chat(
|
||
client=client,
|
||
model_name=model_name,
|
||
msg=context_msg + [UserMessage(content=usermsg)] + tool_msg,
|
||
tools=tools.get_tools_list()
|
||
)
|
||
choice = response.choices[0]
|
||
context.append(
|
||
UserMessage(content=usermsg).as_dict(), target.id, target.private
|
||
)
|
||
# context.append(tool_msg, target.id, target.private)
|
||
context.append(choice.message.as_dict(), target.id, target.private)
|
||
await UniMessage(str(choice.message.content)).send(reply_to=True)
|
||
except Exception as e:
|
||
await UniMessage(str(e) + suggest_solution(str(e))).send()
|
||
traceback.print_exc()
|
||
return
|
||
|
||
|
||
with contextlib.suppress(ImportError): # 优化先不做（）
|
||
import nonebot.adapters.onebot.v11 # type: ignore
|
||
from .azure_onebot import poke_notify
|
||
|
||
|
||
@poke_notify.handle()
|
||
async def poke(event: Event):
|
||
|
||
user_id = event.get_user_id()
|
||
nicknames = await get_nicknames()
|
||
user_nickname = nicknames.get(user_id, "")
|
||
try:
|
||
if config.marshoai_poke_suffix != "":
|
||
response = await make_chat(
|
||
client=client,
|
||
model_name=model_name,
|
||
msg=[
|
||
get_prompt(),
|
||
UserMessage(
|
||
content=f"*{user_nickname}{config.marshoai_poke_suffix}"
|
||
),
|
||
],
|
||
)
|
||
choice = response.choices[0]
|
||
if choice["finish_reason"] == CompletionsFinishReason.STOPPED:
|
||
await UniMessage(" " + str(choice.message.content)).send(
|
||
at_sender=True
|
||
)
|
||
except Exception as e:
|
||
await UniMessage(str(e) + suggest_solution(str(e))).send()
|
||
traceback.print_exc()
|
||
return
|
||
|
||
|
||
@driver.on_shutdown
|
||
async def auto_backup_context():
|
||
for target_info in target_list:
|
||
target_id, target_private = target_info
|
||
contexts_data = context.build(target_id, target_private)
|
||
if target_private:
|
||
target_uid = "private_" + target_id
|
||
else:
|
||
target_uid = "group_" + target_id
|
||
await save_context_to_json(f"back_up_context_{target_uid}", contexts_data, "contexts/backup")
|
||
logger.info(f"已保存会话 {target_id} 的上下文备份，将在下次对话时恢复~")
|