kotones-auto-assistant/kotonebot/backend/debug/server.py

226 lines
7.2 KiB
Python

import time
import asyncio
import inspect
import threading
import traceback
import subprocess
from io import StringIO
from pathlib import Path
from typing import Literal
from collections import deque
from contextlib import redirect_stdout
import cv2
import uvicorn
from thefuzz import fuzz
from pydantic import BaseModel
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, Response
from fastapi import FastAPI, WebSocket, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import kotonebot
import kotonebot.backend
import kotonebot.backend.context
from kotonebot.backend.core import HintBox, Image
from ..context import manual_context
from . import vars as debug_vars
from .vars import WSImage, WSMessageData, WSMessage, WSCallstack
app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
# 获取当前文件夹路径
CURRENT_DIR = Path(__file__).parent
APP_DIR = Path.cwd()
class File(BaseModel):
name: str
full_path: str
type: Literal["file", "dir"]
@app.get("/api/read_file")
async def read_file(path: str):
"""读取文件内容"""
try:
# 确保路径在当前目录下
full_path = (APP_DIR / path).resolve()
if not Path(full_path).is_relative_to(APP_DIR):
raise HTTPException(status_code=403, detail="Access denied")
if not full_path.exists():
raise HTTPException(status_code=404, detail="File not found")
# 添加缓存控制头
headers = {
"Cache-Control": "public, max-age=3600", # 缓存1小时
"ETag": f'"{hash(full_path)}"' # 使用full_path的哈希值作为ETag
}
return FileResponse(full_path, headers=headers)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/read_memory")
async def read_memory(key: str):
"""读取内存中的数据"""
try:
image = None
if (image := debug_vars._read_image(key)) is not None:
pass
else:
raise HTTPException(status_code=404, detail="Key not found")
# 编码图片
encode_params = [cv2.IMWRITE_PNG_COMPRESSION, 4]
_, buffer = cv2.imencode('.png', image, encode_params)
# 添加缓存控制头
headers = {
"Cache-Control": "public, max-age=3600", # 缓存1小时
"ETag": f'"{hash(key)}"' # 使用key的哈希值作为ETag
}
return Response(
buffer.tobytes(),
media_type="image/jpeg",
headers=headers
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/screenshot")
def screenshot():
from ..context import device
img = device.screenshot()
buff = cv2.imencode('.png', img)[1].tobytes()
return Response(buff, media_type="image/png")
class RunCodeRequest(BaseModel):
code: str
@app.post("/api/code/run")
async def run_code(request: RunCodeRequest):
event = asyncio.Event()
stdout = StringIO()
code = f"from kotonebot import *\n" + request.code
result = {}
def _runner():
nonlocal result
from kotonebot.backend.context import vars as context_vars
try:
with manual_context():
global_vars = dict(vars(kotonebot.backend.context))
with redirect_stdout(stdout):
exec(code, global_vars)
result = {"status": "ok", "result": stdout.getvalue()}
except (Exception) as e:
result = {"status": "error", "result": stdout.getvalue(), "message": str(e), "traceback": traceback.format_exc()}
except KeyboardInterrupt as e:
result = {"status": "error", "result": stdout.getvalue(), "message": str(e), "traceback": traceback.format_exc()}
finally:
context_vars.interrupted.clear()
event.set()
threading.Thread(target=_runner, daemon=True).start()
await event.wait()
return result
@app.get("/api/code/stop")
async def stop_code():
from kotonebot.backend.context import vars
vars.interrupted.set()
while vars.interrupted.is_set():
await asyncio.sleep(0.1)
return {"status": "ok"}
@app.get("/api/fs/list_dir")
def list_dir(path: str) -> list[File]:
result = []
for item in Path(path).iterdir():
result.append(File(
name=item.name,
full_path=str(item),
type="file" if item.is_file() else "dir"
))
return result
@app.get("/api/resources/autocomplete")
def autocomplete(class_path: str) -> list[str]:
from kotonebot.kaa import R # HACK: hardcode
class_names = class_path.split(".")[:-1]
target_class = R
# 定位到目标类
for name in class_names:
target_class = getattr(target_class, name, None)
if target_class is None:
return []
# 获取目标类的所有属性
attrs = [attr for attr in dir(target_class) if not attr.startswith("_")]
filtered_attrs = []
for attr in attrs:
if inspect.isclass(getattr(target_class, attr)):
filtered_attrs.append(attr)
elif isinstance(getattr(target_class, attr), (Image, HintBox)):
filtered_attrs.append(attr)
attrs = filtered_attrs
# 排序
attrs.sort(key=lambda x: fuzz.ratio(x, class_path), reverse=True)
return attrs
@app.get("/api/ping")
async def ping():
return {"status": "ok"}
message_queue = deque()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
try:
while True:
if len(message_queue) > 0:
message = message_queue.popleft()
await websocket.send_json(message)
await asyncio.sleep(0.1)
except:
await websocket.close()
def send_ws_message(title: str, image: list[str], text: str = '', callstack: list[WSCallstack] = [], wait: bool = False):
"""发送 WebSocket 消息"""
message = WSMessage(
type="visual",
data=WSMessageData(
image=WSImage(type="memory", value=image),
name=title,
details=text,
timestamp=int(time.time() * 1000),
callstack=callstack
)
)
message_queue.append(message.dict())
if wait:
while len(message_queue) > 0:
time.sleep(0.3)
thread = None
def start_server():
global thread
def run_server():
uvicorn.run(app, host="127.0.0.1", port=8000, log_level='critical' if debug_vars.debug.hide_server_log else None)
if thread is None:
thread = threading.Thread(target=run_server, daemon=True)
thread.start()
def wait_message_all_done():
global thread
def _wait():
while len(message_queue) > 0:
time.sleep(0.1)
if thread is not None:
threading.Thread(target=_wait, daemon=True).start()
if __name__ == "__main__":
from kotonebot.backend.context import init_context
init_context()
debug_vars.debug.hide_server_log = False
process = subprocess.Popen(["pylsp", "--port", "5479", "--ws"])
print("LSP started. PID=", process.pid)
uvicorn.run(app, host="127.0.0.1", port=8000, log_level='critical' if debug_vars.debug.hide_server_log else None)
process.kill()