mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-20 02:55:08 +08:00
263 lines
10 KiB
Python
263 lines
10 KiB
Python
import uuid
|
|
import json
|
|
import os
|
|
from .route import Route, Response, RouteContext
|
|
from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queue_mgr
|
|
from quart import request, Response as QuartResponse, g, make_response
|
|
from astrbot.core.db import BaseDatabase
|
|
import asyncio
|
|
from astrbot.core import logger
|
|
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
|
|
from astrbot.core.utils.astrbot_path import get_astrbot_data_path
|
|
|
|
|
|
class ChatRoute(Route):
|
|
def __init__(
|
|
self,
|
|
context: RouteContext,
|
|
db: BaseDatabase,
|
|
core_lifecycle: AstrBotCoreLifecycle,
|
|
) -> None:
|
|
super().__init__(context)
|
|
self.routes = {
|
|
"/chat/send": ("POST", self.chat),
|
|
"/chat/new_conversation": ("GET", self.new_conversation),
|
|
"/chat/conversations": ("GET", self.get_conversations),
|
|
"/chat/get_conversation": ("GET", self.get_conversation),
|
|
"/chat/delete_conversation": ("GET", self.delete_conversation),
|
|
"/chat/rename_conversation": ("POST", self.rename_conversation),
|
|
"/chat/get_file": ("GET", self.get_file),
|
|
"/chat/post_image": ("POST", self.post_image),
|
|
"/chat/post_file": ("POST", self.post_file),
|
|
"/chat/status": ("GET", self.status),
|
|
}
|
|
self.db = db
|
|
self.core_lifecycle = core_lifecycle
|
|
self.register_routes()
|
|
self.imgs_dir = os.path.join(get_astrbot_data_path(), "webchat", "imgs")
|
|
os.makedirs(self.imgs_dir, exist_ok=True)
|
|
|
|
self.supported_imgs = ["jpg", "jpeg", "png", "gif", "webp"]
|
|
|
|
async def status(self):
|
|
has_llm_enabled = (
|
|
self.core_lifecycle.provider_manager.curr_provider_inst is not None
|
|
)
|
|
has_stt_enabled = (
|
|
self.core_lifecycle.provider_manager.curr_stt_provider_inst is not None
|
|
)
|
|
return (
|
|
Response()
|
|
.ok(data={"llm_enabled": has_llm_enabled, "stt_enabled": has_stt_enabled})
|
|
.__dict__
|
|
)
|
|
|
|
async def get_file(self):
|
|
filename = request.args.get("filename")
|
|
if not filename:
|
|
return Response().error("Missing key: filename").__dict__
|
|
|
|
try:
|
|
file_path = os.path.join(self.imgs_dir, os.path.basename(filename))
|
|
real_file_path = os.path.realpath(file_path)
|
|
real_imgs_dir = os.path.realpath(self.imgs_dir)
|
|
|
|
if not real_file_path.startswith(real_imgs_dir):
|
|
return Response().error("Invalid file path").__dict__
|
|
|
|
with open(real_file_path, "rb") as f:
|
|
filename_ext = os.path.splitext(filename)[1].lower()
|
|
|
|
if filename_ext == ".wav":
|
|
return QuartResponse(f.read(), mimetype="audio/wav")
|
|
elif filename_ext[1:] in self.supported_imgs:
|
|
return QuartResponse(f.read(), mimetype="image/jpeg")
|
|
else:
|
|
return QuartResponse(f.read())
|
|
|
|
except (FileNotFoundError, OSError):
|
|
return Response().error("File access error").__dict__
|
|
|
|
async def post_image(self):
|
|
post_data = await request.files
|
|
if "file" not in post_data:
|
|
return Response().error("Missing key: file").__dict__
|
|
|
|
file = post_data["file"]
|
|
filename = str(uuid.uuid4()) + ".jpg"
|
|
path = os.path.join(self.imgs_dir, filename)
|
|
await file.save(path)
|
|
|
|
return Response().ok(data={"filename": filename}).__dict__
|
|
|
|
async def post_file(self):
|
|
post_data = await request.files
|
|
if "file" not in post_data:
|
|
return Response().error("Missing key: file").__dict__
|
|
|
|
file = post_data["file"]
|
|
filename = f"{str(uuid.uuid4())}"
|
|
# 通过文件格式判断文件类型
|
|
if file.content_type.startswith("audio"):
|
|
filename += ".wav"
|
|
|
|
path = os.path.join(self.imgs_dir, filename)
|
|
await file.save(path)
|
|
|
|
return Response().ok(data={"filename": filename}).__dict__
|
|
|
|
async def chat(self):
|
|
username = g.get("username", "guest")
|
|
|
|
post_data = await request.json
|
|
if "message" not in post_data and "image_url" not in post_data:
|
|
return Response().error("Missing key: message or image_url").__dict__
|
|
|
|
if "conversation_id" not in post_data:
|
|
return Response().error("Missing key: conversation_id").__dict__
|
|
|
|
message = post_data["message"]
|
|
conversation_id = post_data["conversation_id"]
|
|
image_url = post_data.get("image_url")
|
|
audio_url = post_data.get("audio_url")
|
|
selected_provider = post_data.get("selected_provider")
|
|
selected_model = post_data.get("selected_model")
|
|
if not message and not image_url and not audio_url:
|
|
return (
|
|
Response()
|
|
.error("Message and image_url and audio_url are empty")
|
|
.__dict__
|
|
)
|
|
if not conversation_id:
|
|
return Response().error("conversation_id is empty").__dict__
|
|
|
|
# Get conversation-specific queues
|
|
back_queue = webchat_queue_mgr.get_or_create_back_queue(conversation_id)
|
|
|
|
# append user message
|
|
conversation = self.db.get_conversation_by_user_id(username, conversation_id)
|
|
try:
|
|
history = json.loads(conversation.history)
|
|
except BaseException as e:
|
|
logger.error(f"Failed to parse conversation history: {e}")
|
|
history = []
|
|
new_his = {"type": "user", "message": message}
|
|
if image_url:
|
|
new_his["image_url"] = image_url
|
|
if audio_url:
|
|
new_his["audio_url"] = audio_url
|
|
history.append(new_his)
|
|
self.db.update_conversation(
|
|
username, conversation_id, history=json.dumps(history)
|
|
)
|
|
|
|
async def stream():
|
|
try:
|
|
while True:
|
|
try:
|
|
result = await asyncio.wait_for(back_queue.get(), timeout=10)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
if not result:
|
|
continue
|
|
|
|
result_text = result["data"]
|
|
type = result.get("type")
|
|
cid = result.get("cid")
|
|
streaming = result.get("streaming", False)
|
|
yield f"data: {json.dumps(result, ensure_ascii=False)}\n\n"
|
|
await asyncio.sleep(0.05)
|
|
|
|
if type == "end":
|
|
break
|
|
elif (streaming and type == "complete") or not streaming:
|
|
# append bot message
|
|
conversation = self.db.get_conversation_by_user_id(
|
|
username, cid
|
|
)
|
|
try:
|
|
history = json.loads(conversation.history)
|
|
except BaseException as e:
|
|
logger.error(f"Failed to parse conversation history: {e}")
|
|
history = []
|
|
history.append({"type": "bot", "message": result_text})
|
|
self.db.update_conversation(
|
|
username, cid, history=json.dumps(history)
|
|
)
|
|
|
|
except BaseException as _:
|
|
logger.debug(f"用户 {username} 断开聊天长连接。")
|
|
return
|
|
|
|
# Put message to conversation-specific queue
|
|
chat_queue = webchat_queue_mgr.get_or_create_queue(conversation_id)
|
|
await chat_queue.put(
|
|
(
|
|
username,
|
|
conversation_id,
|
|
{
|
|
"message": message,
|
|
"image_url": image_url, # list
|
|
"audio_url": audio_url,
|
|
"selected_provider": selected_provider,
|
|
"selected_model": selected_model,
|
|
},
|
|
)
|
|
)
|
|
|
|
response = await make_response(
|
|
stream(),
|
|
{
|
|
"Content-Type": "text/event-stream",
|
|
"Cache-Control": "no-cache",
|
|
"Transfer-Encoding": "chunked",
|
|
"Connection": "keep-alive",
|
|
},
|
|
)
|
|
return response
|
|
|
|
async def delete_conversation(self):
|
|
username = g.get("username", "guest")
|
|
conversation_id = request.args.get("conversation_id")
|
|
if not conversation_id:
|
|
return Response().error("Missing key: conversation_id").__dict__
|
|
|
|
# Clean up queues when deleting conversation
|
|
webchat_queue_mgr.remove_queues(conversation_id)
|
|
self.db.delete_conversation(username, conversation_id)
|
|
return Response().ok().__dict__
|
|
|
|
async def new_conversation(self):
|
|
username = g.get("username", "guest")
|
|
conversation_id = str(uuid.uuid4())
|
|
self.db.new_conversation(username, conversation_id)
|
|
return Response().ok(data={"conversation_id": conversation_id}).__dict__
|
|
|
|
async def rename_conversation(self):
|
|
username = g.get("username", "guest")
|
|
post_data = await request.json
|
|
if "conversation_id" not in post_data or "title" not in post_data:
|
|
return Response().error("Missing key: conversation_id or title").__dict__
|
|
|
|
conversation_id = post_data["conversation_id"]
|
|
title = post_data["title"]
|
|
|
|
self.db.update_conversation_title(username, conversation_id, title=title)
|
|
return Response().ok(message="重命名成功！").__dict__
|
|
|
|
async def get_conversations(self):
|
|
username = g.get("username", "guest")
|
|
conversations = self.db.get_conversations(username)
|
|
return Response().ok(data=conversations).__dict__
|
|
|
|
async def get_conversation(self):
|
|
username = g.get("username", "guest")
|
|
conversation_id = request.args.get("conversation_id")
|
|
if not conversation_id:
|
|
return Response().error("Missing key: conversation_id").__dict__
|
|
|
|
conversation = self.db.get_conversation_by_user_id(username, conversation_id)
|
|
|
|
return Response().ok(data=conversation).__dict__
|