diff --git a/AIGC/main.py b/AIGC/main.py index b754423..fe84656 100644 --- a/AIGC/main.py +++ b/AIGC/main.py @@ -15,7 +15,7 @@ from Utils.AIGCLog import AIGCLog from ollama import Client, ResponseError app = FastAPI(title = "AI 通信服务") -logger = AIGCLog(name = "AIGC", log_file = "./Log/aigc.log") +logger = AIGCLog(name = "AIGC", log_file = "aigc.log") parser = argparse.ArgumentParser() parser.add_argument('--model', type=str, default='deepseek-r1:1.5b', @@ -31,12 +31,13 @@ lastPrompt = "" ollamaClient = Client(host='http://localhost:11434') async def heartbeat(websocket: WebSocket): - while True: - await asyncio.sleep(30) # 每30秒发送一次心跳 - try: - await senddata(websocket, 0, []) - except ConnectionClosed: - break # 连接已关闭时退出循环 + pass + # while True: + # await asyncio.sleep(30) # 每30秒发送一次心跳 + # try: + # await senddata(websocket, 0, []) + # except ConnectionClosed: + # break # 连接已关闭时退出循环 #statuscode -1 服务器运行错误 0 心跳标志 1 正常 2 输出异常 async def senddata(websocket: WebSocket, statusCode: int, messages: List[str]): @@ -63,6 +64,7 @@ async def websocket_endpoint(websocket: WebSocket, client_id: str): data = await websocket.receive_text() logger.log(logging.INFO, f"收到UE5消息 [{client_id}]: {data}") success, prompt = process_prompt(data) + global lastPrompt lastPrompt = prompt # 调用AI生成响应 if(success): @@ -175,7 +177,7 @@ async def generateAIChat(prompt: str, websocket: WebSocket): match = re.search(pattern, think_remove_text, re.DOTALL) if not match: - if await reGenerateAIChat(prompt, websocket): + if await reGenerateAIChat(lastPrompt, websocket): pass else: logger.log(logging.ERROR, "请更换prompt,或者升级模型大小") @@ -183,9 +185,9 @@ async def generateAIChat(prompt: str, websocket: WebSocket): else: core_dialog = match.group(1).strip() - dialog_lines = core_dialog.split('\n') + dialog_lines = [line for line in core_dialog.split('\n') if line.strip()] if len(dialog_lines) != 4: - if await reGenerateAIChat(prompt, websocket): + if await reGenerateAIChat(lastPrompt, websocket): pass else: logger.log(logging.ERROR, "请更换prompt,或者升级模型大小") @@ -204,6 +206,7 @@ async def reGenerateAIChat(prompt: str, websocket: WebSocket): logger.log(logging.ERROR, f"AI输出格式不正确,重新进行生成 {regenerateCount}/{maxAIRegerateCount}") await senddata(websocket, 2, messages=["ai生成格式不正确, 重新进行生成"]) await asyncio.sleep(0) + prompt = prompt + "补充:上一次的输出格式错误,严格执行prompt中第5条的输出格式要求" await generateAIChat(prompt, websocket) return True else: