Compare commits
2
Commits
9e2667d604
...
2982bc30a4
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2982bc30a4 | ||
|
|
aaeac343d3 |
+187
-10
@@ -9,11 +9,21 @@ import sys
|
|||||||
import os
|
import os
|
||||||
import threading
|
import threading
|
||||||
import hashlib
|
import hashlib
|
||||||
|
from collections import OrderedDict
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Optional, Dict, Any, List
|
from typing import Optional, Dict, Any, List
|
||||||
import websockets
|
import websockets
|
||||||
from colorama import init, Fore, Style as ColoramaStyle
|
from colorama import init, Fore, Style as ColoramaStyle
|
||||||
|
|
||||||
|
# rich 为可选依赖:有则流式渲染 Markdown,无则回退纯文本流式输出
|
||||||
|
try:
|
||||||
|
from rich.console import Console
|
||||||
|
from rich.live import Live
|
||||||
|
from rich.markdown import Markdown
|
||||||
|
_RICH = True
|
||||||
|
except ImportError:
|
||||||
|
_RICH = False
|
||||||
|
|
||||||
init(autoreset=True)
|
init(autoreset=True)
|
||||||
|
|
||||||
class Colors:
|
class Colors:
|
||||||
@@ -75,10 +85,15 @@ class WebSocketClient:
|
|||||||
except json.JSONDecodeError:
|
except json.JSONDecodeError:
|
||||||
self.callback("chunk", message)
|
self.callback("chunk", message)
|
||||||
except websockets.exceptions.ConnectionClosed:
|
except websockets.exceptions.ConnectionClosed:
|
||||||
self.connected = False
|
pass
|
||||||
self.callback("system", "❌ 连接已断开")
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.callback("error", f"接收消息错误: {e}")
|
self.callback("error", f"接收消息错误: {e}")
|
||||||
|
finally:
|
||||||
|
# 统一在此通知断线;手动 disconnect() 已把 connected 置 False,不会误触发自动重连
|
||||||
|
was_connected = self.connected
|
||||||
|
self.connected = False
|
||||||
|
if was_connected:
|
||||||
|
self.callback("disconnected", "❌ 连接已断开")
|
||||||
|
|
||||||
async def _handle_message(self, data: dict):
|
async def _handle_message(self, data: dict):
|
||||||
msg_type = data.get('type', '')
|
msg_type = data.get('type', '')
|
||||||
@@ -134,13 +149,27 @@ class SimpleCLI:
|
|||||||
self.ai_full = ""
|
self.ai_full = ""
|
||||||
self.ai_prefix_shown = False
|
self.ai_prefix_shown = False
|
||||||
|
|
||||||
# 已显示消息的去重
|
# 已显示消息的去重(LRU 上限,防长对话内存无限增长)
|
||||||
self.displayed_messages: set = set()
|
self.DISPLAYED_MAX = 500
|
||||||
|
self.displayed_messages: OrderedDict = OrderedDict()
|
||||||
|
|
||||||
# 锁和状态
|
# 锁和状态
|
||||||
self.lock = threading.Lock()
|
self.lock = threading.Lock()
|
||||||
self.current_input = ""
|
self.current_input = ""
|
||||||
|
|
||||||
|
# 标准输入:单守护线程阻塞读 + asyncio.Queue 中转(避免 Ctrl+C 泄漏线程)
|
||||||
|
self._input_queue: asyncio.Queue = asyncio.Queue()
|
||||||
|
self._stdin_thread: Optional[threading.Thread] = None
|
||||||
|
|
||||||
|
# 自动重连状态
|
||||||
|
self._manual_disconnect = False # 用户主动 /disconnect、/quit、/connect 后置 True,抑制自动重连
|
||||||
|
self._reconnect_task: Optional[asyncio.Task] = None
|
||||||
|
|
||||||
|
# Markdown 流式渲染状态(rich 可用时启用)
|
||||||
|
self._console = Console() if _RICH else None
|
||||||
|
self._md_live = None # rich.Live 实例(渲染进行中)
|
||||||
|
self._md_buffer = "" # 当前累计的 AI 回复全文
|
||||||
|
|
||||||
# 流式输出状态
|
# 流式输出状态
|
||||||
self.streaming_line = False # 是否正在流式输出行
|
self.streaming_line = False # 是否正在流式输出行
|
||||||
self.in_tool_phase = False # 是否在工具调用阶段
|
self.in_tool_phase = False # 是否在工具调用阶段
|
||||||
@@ -165,6 +194,7 @@ class SimpleCLI:
|
|||||||
|
|
||||||
def _print_line(self, line: str = "", end: str = "\n"):
|
def _print_line(self, line: str = "", end: str = "\n"):
|
||||||
"""打印一行,并恢复提示符"""
|
"""打印一行,并恢复提示符"""
|
||||||
|
self._md_stream_stop() # 打印普通行前,先定格进行中的 Markdown 渲染
|
||||||
with self.lock:
|
with self.lock:
|
||||||
current = self.current_input
|
current = self.current_input
|
||||||
sys.stdout.write('\r' + ' ' * 120 + '\r')
|
sys.stdout.write('\r' + ' ' * 120 + '\r')
|
||||||
@@ -242,6 +272,17 @@ class SimpleCLI:
|
|||||||
if key in self.displayed_messages:
|
if key in self.displayed_messages:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
# rich 可用:累计文本直接整体渲染为 Markdown(服务端 chunk 为累计全文)
|
||||||
|
if _RICH:
|
||||||
|
with self.lock:
|
||||||
|
self.waiting_for_response = True
|
||||||
|
self.prompt_shown = False
|
||||||
|
self.streaming_line = True
|
||||||
|
self._md_stream_update(clean)
|
||||||
|
self.ai_previous = clean
|
||||||
|
self.ai_full = clean
|
||||||
|
return
|
||||||
|
|
||||||
# 检测增量
|
# 检测增量
|
||||||
if self.ai_previous and clean.startswith(self.ai_previous):
|
if self.ai_previous and clean.startswith(self.ai_previous):
|
||||||
new_part = clean[len(self.ai_previous):]
|
new_part = clean[len(self.ai_previous):]
|
||||||
@@ -275,6 +316,44 @@ class SimpleCLI:
|
|||||||
self.ai_previous = clean
|
self.ai_previous = clean
|
||||||
self.ai_full = clean
|
self.ai_full = clean
|
||||||
|
|
||||||
|
# ── Markdown 流式渲染(rich)─────────────────────────────────────────
|
||||||
|
def _md_stream_start(self):
|
||||||
|
"""开启 Live 渲染区域(若 rich 不可用或已开启则忽略)。"""
|
||||||
|
if not _RICH or self._md_live is not None:
|
||||||
|
return
|
||||||
|
self._md_buffer = ""
|
||||||
|
# transient=False:结束后保留最终渲染结果;vertical_overflow=visible 不裁剪长文
|
||||||
|
self._md_live = Live(Markdown(""), console=self._console,
|
||||||
|
refresh_per_second=8, vertical_overflow="visible")
|
||||||
|
self._md_live.start()
|
||||||
|
|
||||||
|
def _md_stream_update(self, full_text: str):
|
||||||
|
"""用累计全文刷新渲染(服务端 chunk 即累计文本)。"""
|
||||||
|
if not _RICH:
|
||||||
|
return
|
||||||
|
self._md_stream_start()
|
||||||
|
self._md_buffer = full_text
|
||||||
|
self._md_live.update(Markdown(full_text))
|
||||||
|
|
||||||
|
def _md_stream_stop(self):
|
||||||
|
"""结束渲染:最终刷新一次后停住(内容定格保留在终端)。"""
|
||||||
|
if self._md_live is not None:
|
||||||
|
try:
|
||||||
|
if self._md_buffer:
|
||||||
|
self._md_live.update(Markdown(self._md_buffer))
|
||||||
|
self._md_live.stop()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
self._md_live = None
|
||||||
|
|
||||||
|
def _remember_displayed(self, key: str):
|
||||||
|
"""记录已显示消息(LRU:超 DISPLAYED_MAX 淘汰最旧,防长对话内存膨胀)。"""
|
||||||
|
self.displayed_messages[key] = None
|
||||||
|
self.displayed_messages.move_to_end(key)
|
||||||
|
if len(self.displayed_messages) > self.DISPLAYED_MAX:
|
||||||
|
self.displayed_messages.popitem(last=False)
|
||||||
|
|
||||||
def display(self, msg_type: str, content: str, extra: Any = None):
|
def display(self, msg_type: str, content: str, extra: Any = None):
|
||||||
"""显示消息"""
|
"""显示消息"""
|
||||||
if msg_type == 'chunk':
|
if msg_type == 'chunk':
|
||||||
@@ -286,6 +365,7 @@ class SimpleCLI:
|
|||||||
elif msg_type == 'task_complete':
|
elif msg_type == 'task_complete':
|
||||||
self._handle_task_complete(extra)
|
self._handle_task_complete(extra)
|
||||||
elif msg_type == 'dialogue_ended':
|
elif msg_type == 'dialogue_ended':
|
||||||
|
self._md_stream_stop()
|
||||||
with self.lock:
|
with self.lock:
|
||||||
if self.streaming_line:
|
if self.streaming_line:
|
||||||
sys.stdout.write('\n')
|
sys.stdout.write('\n')
|
||||||
@@ -302,6 +382,11 @@ class SimpleCLI:
|
|||||||
elif msg_type == 'system':
|
elif msg_type == 'system':
|
||||||
self._finish_stream()
|
self._finish_stream()
|
||||||
self._print_line(f"{Colors.SYSTEM}{content}{Colors.RESET}")
|
self._print_line(f"{Colors.SYSTEM}{content}{Colors.RESET}")
|
||||||
|
elif msg_type == 'disconnected':
|
||||||
|
self.connected = False
|
||||||
|
self._finish_stream()
|
||||||
|
self._print_line(f"{Colors.SYSTEM}{content}{Colors.RESET}")
|
||||||
|
self._schedule_reconnect()
|
||||||
elif msg_type == 'info':
|
elif msg_type == 'info':
|
||||||
self._finish_stream()
|
self._finish_stream()
|
||||||
self._print_line(f"{Colors.INFO}ℹ️ {content}{Colors.RESET}")
|
self._print_line(f"{Colors.INFO}ℹ️ {content}{Colors.RESET}")
|
||||||
@@ -313,6 +398,7 @@ class SimpleCLI:
|
|||||||
"""处理流式块"""
|
"""处理流式块"""
|
||||||
# 检查是否是工具调用相关的消息
|
# 检查是否是工具调用相关的消息
|
||||||
if content.startswith('🔧 调用工具'):
|
if content.startswith('🔧 调用工具'):
|
||||||
|
self._md_stream_stop() # 工具消息与 Markdown 渲染区域交替,先定格
|
||||||
with self.lock:
|
with self.lock:
|
||||||
self.in_tool_phase = True
|
self.in_tool_phase = True
|
||||||
self.prompt_shown = False
|
self.prompt_shown = False
|
||||||
@@ -326,6 +412,7 @@ class SimpleCLI:
|
|||||||
|
|
||||||
# 检查是否是工具结果
|
# 检查是否是工具结果
|
||||||
if content.startswith(('✅ 对话记忆已保存', '💾 保存记忆', '❌', '⚠️')):
|
if content.startswith(('✅ 对话记忆已保存', '💾 保存记忆', '❌', '⚠️')):
|
||||||
|
self._md_stream_stop()
|
||||||
with self.lock:
|
with self.lock:
|
||||||
# 打印工具结果
|
# 打印工具结果
|
||||||
sys.stdout.write(f"{Colors.TASK}{content}{Colors.RESET}\n")
|
sys.stdout.write(f"{Colors.TASK}{content}{Colors.RESET}\n")
|
||||||
@@ -350,6 +437,27 @@ class SimpleCLI:
|
|||||||
|
|
||||||
def _handle_complete(self, content: str):
|
def _handle_complete(self, content: str):
|
||||||
"""处理完成事件 - 只负责结束流式输出,不重复显示内容"""
|
"""处理完成事件 - 只负责结束流式输出,不重复显示内容"""
|
||||||
|
# rich 渲染路径:最终全文已在 Live 区域定格,只需收尾并恢复提示符
|
||||||
|
if _RICH:
|
||||||
|
clean = self._clean_content(content) if content else ""
|
||||||
|
if clean:
|
||||||
|
self._md_stream_update(clean)
|
||||||
|
self._remember_displayed(clean[:200])
|
||||||
|
self._md_stream_stop()
|
||||||
|
with self.lock:
|
||||||
|
self.in_tool_phase = False
|
||||||
|
self.waiting_for_response = False
|
||||||
|
self.streaming_line = False
|
||||||
|
self.prompt_shown = True
|
||||||
|
current = self.current_input
|
||||||
|
sys.stdout.write('\r' + ' ' * 120 + '\r')
|
||||||
|
sys.stdout.write(self._get_prompt() + current)
|
||||||
|
sys.stdout.flush()
|
||||||
|
self.ai_previous = ""
|
||||||
|
self.ai_full = ""
|
||||||
|
self.ai_prefix_shown = False
|
||||||
|
return
|
||||||
|
|
||||||
# 如果之前在工具阶段,现在退出
|
# 如果之前在工具阶段,现在退出
|
||||||
if self.in_tool_phase:
|
if self.in_tool_phase:
|
||||||
with self.lock:
|
with self.lock:
|
||||||
@@ -406,7 +514,7 @@ class SimpleCLI:
|
|||||||
return
|
return
|
||||||
|
|
||||||
# 只有当内容没有被流式显示时才显示完整消息
|
# 只有当内容没有被流式显示时才显示完整消息
|
||||||
self.displayed_messages.add(key)
|
self._remember_displayed(key)
|
||||||
# 如果之前没有流式输出,现在显示完整消息
|
# 如果之前没有流式输出,现在显示完整消息
|
||||||
self._print_line(f"{Colors.AI}🤖 AI助手: {clean}")
|
self._print_line(f"{Colors.AI}🤖 AI助手: {clean}")
|
||||||
|
|
||||||
@@ -466,7 +574,7 @@ class SimpleCLI:
|
|||||||
clean = content.strip()
|
clean = content.strip()
|
||||||
key = clean[:200]
|
key = clean[:200]
|
||||||
if key not in self.displayed_messages:
|
if key not in self.displayed_messages:
|
||||||
self.displayed_messages.add(key)
|
self._remember_displayed(key)
|
||||||
self._print_line(f"{Colors.TASK}⚡ 任务 - {task_title}: {clean}")
|
self._print_line(f"{Colors.TASK}⚡ 任务 - {task_title}: {clean}")
|
||||||
del self._task_buffers[task_title]
|
del self._task_buffers[task_title]
|
||||||
self.streaming_line = False
|
self.streaming_line = False
|
||||||
@@ -519,10 +627,15 @@ class SimpleCLI:
|
|||||||
self.display("error", "用法: /connect ws://host:port")
|
self.display("error", "用法: /connect ws://host:port")
|
||||||
return
|
return
|
||||||
uri = parts[1]
|
uri = parts[1]
|
||||||
|
# 切换服务器:抑制旧连接的自动重连
|
||||||
|
self._manual_disconnect = True
|
||||||
|
if self._reconnect_task and not self._reconnect_task.done():
|
||||||
|
self._reconnect_task.cancel()
|
||||||
if self.client and self.client.connected:
|
if self.client and self.client.connected:
|
||||||
await self.client.disconnect()
|
await self.client.disconnect()
|
||||||
self.client = WebSocketClient(uri, self.display)
|
self.client = WebSocketClient(uri, self.display)
|
||||||
if await self.client.connect():
|
if await self.client.connect():
|
||||||
|
self._manual_disconnect = False
|
||||||
self.connected = True
|
self.connected = True
|
||||||
self.display("system", f"✅ 已连接到 {uri}")
|
self.display("system", f"✅ 已连接到 {uri}")
|
||||||
asyncio.create_task(self.client.receive_messages())
|
asyncio.create_task(self.client.receive_messages())
|
||||||
@@ -530,6 +643,9 @@ class SimpleCLI:
|
|||||||
self.client = None
|
self.client = None
|
||||||
self.connected = False
|
self.connected = False
|
||||||
elif cmd == '/disconnect':
|
elif cmd == '/disconnect':
|
||||||
|
self._manual_disconnect = True # 主动断开,不自动重连
|
||||||
|
if self._reconnect_task and not self._reconnect_task.done():
|
||||||
|
self._reconnect_task.cancel()
|
||||||
if self.client:
|
if self.client:
|
||||||
await self.client.disconnect()
|
await self.client.disconnect()
|
||||||
self.client = None
|
self.client = None
|
||||||
@@ -543,6 +659,7 @@ class SimpleCLI:
|
|||||||
else:
|
else:
|
||||||
self.display("system", "未连接")
|
self.display("system", "未连接")
|
||||||
elif cmd == '/clear':
|
elif cmd == '/clear':
|
||||||
|
self._md_stream_stop()
|
||||||
os.system('clear' if os.name == 'posix' else 'cls')
|
os.system('clear' if os.name == 'posix' else 'cls')
|
||||||
self.displayed_messages.clear()
|
self.displayed_messages.clear()
|
||||||
self.current_input = ""
|
self.current_input = ""
|
||||||
@@ -556,6 +673,9 @@ class SimpleCLI:
|
|||||||
elif cmd == '/quit':
|
elif cmd == '/quit':
|
||||||
self.display("system", "👋 再见!")
|
self.display("system", "👋 再见!")
|
||||||
self.running = False
|
self.running = False
|
||||||
|
self._manual_disconnect = True
|
||||||
|
if self._reconnect_task and not self._reconnect_task.done():
|
||||||
|
self._reconnect_task.cancel()
|
||||||
if self.client:
|
if self.client:
|
||||||
await self.client.disconnect()
|
await self.client.disconnect()
|
||||||
return
|
return
|
||||||
@@ -569,6 +689,35 @@ class SimpleCLI:
|
|||||||
self.display("error", f"未知命令: {cmd}")
|
self.display("error", f"未知命令: {cmd}")
|
||||||
self._print_prompt(clear_input=True)
|
self._print_prompt(clear_input=True)
|
||||||
|
|
||||||
|
def _schedule_reconnect(self):
|
||||||
|
"""启动自动重连(若未手动断开且当前无进行中的重连任务)。"""
|
||||||
|
if self._manual_disconnect or not self.running or not self.client:
|
||||||
|
return
|
||||||
|
if self._reconnect_task and not self._reconnect_task.done():
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
self._reconnect_task = asyncio.create_task(self._auto_reconnect_loop(self.client))
|
||||||
|
except RuntimeError:
|
||||||
|
pass # 事件循环正在关闭(程序退出中),放弃重连
|
||||||
|
|
||||||
|
async def _auto_reconnect_loop(self, client: WebSocketClient):
|
||||||
|
"""断线自动重连:指数退避 2s→30s 封顶,直到成功、/disconnect 或退出。"""
|
||||||
|
delay = 2
|
||||||
|
while self.running and not self._manual_disconnect and self.client is client:
|
||||||
|
self.display("info", f"🔄 {delay}s 后自动重连 {client.uri} …(/disconnect 取消)")
|
||||||
|
try:
|
||||||
|
await asyncio.sleep(delay)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
return
|
||||||
|
if not self.running or self._manual_disconnect or self.client is not client:
|
||||||
|
return
|
||||||
|
if await client.connect():
|
||||||
|
self.connected = True
|
||||||
|
self.display("system", f"✅ 已自动重连到 {client.uri}")
|
||||||
|
asyncio.create_task(client.receive_messages())
|
||||||
|
return
|
||||||
|
delay = min(delay * 2, 30)
|
||||||
|
|
||||||
async def _send_interrupt(self):
|
async def _send_interrupt(self):
|
||||||
if self.client and self.client.connected:
|
if self.client and self.client.connected:
|
||||||
interrupt_msg = {"type": "interrupt", "content": "用户按下了 Ctrl+C 中断当前操作"}
|
interrupt_msg = {"type": "interrupt", "content": "用户按下了 Ctrl+C 中断当前操作"}
|
||||||
@@ -582,6 +731,7 @@ class SimpleCLI:
|
|||||||
self.display("error", "未连接到服务器")
|
self.display("error", "未连接到服务器")
|
||||||
return
|
return
|
||||||
|
|
||||||
|
self._md_stream_stop() # 发送新消息前定格上一段渲染
|
||||||
# 重置流式状态
|
# 重置流式状态
|
||||||
with self.lock:
|
with self.lock:
|
||||||
if self.streaming_line:
|
if self.streaming_line:
|
||||||
@@ -601,16 +751,43 @@ class SimpleCLI:
|
|||||||
await self.client.send_message(message)
|
await self.client.send_message(message)
|
||||||
# 不在这里显示提示符,等待AI响应结束后再显示
|
# 不在这里显示提示符,等待AI响应结束后再显示
|
||||||
|
|
||||||
|
def _stdin_reader(self, loop: asyncio.AbstractEventLoop):
|
||||||
|
"""独立守护线程:阻塞读 stdin,经 call_soon_threadsafe 投入队列。
|
||||||
|
|
||||||
|
进程生命周期只启动一次——替代原 run_in_executor 每行一个线程的方案,
|
||||||
|
避免 Ctrl+C 取消 await 后旧线程仍阻塞在 readline 上无限泄漏。
|
||||||
|
"""
|
||||||
|
def _put(item):
|
||||||
|
try:
|
||||||
|
loop.call_soon_threadsafe(self._input_queue.put_nowait, item)
|
||||||
|
except RuntimeError:
|
||||||
|
pass # 事件循环已关闭(程序退出中)
|
||||||
|
|
||||||
|
while self.running:
|
||||||
|
try:
|
||||||
|
line = sys.stdin.readline()
|
||||||
|
except Exception:
|
||||||
|
break
|
||||||
|
if not line: # EOF (Ctrl+D)
|
||||||
|
_put(None)
|
||||||
|
break
|
||||||
|
_put(line)
|
||||||
|
|
||||||
async def run(self):
|
async def run(self):
|
||||||
self.print_banner()
|
self.print_banner()
|
||||||
self._task_buffers = {}
|
self._task_buffers = {}
|
||||||
|
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
if self._stdin_thread is None:
|
||||||
|
self._stdin_thread = threading.Thread(
|
||||||
|
target=self._stdin_reader, args=(loop,),
|
||||||
|
daemon=True, name="stdin-reader")
|
||||||
|
self._stdin_thread.start()
|
||||||
|
|
||||||
while self.running:
|
while self.running:
|
||||||
try:
|
try:
|
||||||
line = await asyncio.get_event_loop().run_in_executor(
|
line = await self._input_queue.get()
|
||||||
None, lambda: sys.stdin.readline()
|
if line is None: # EOF (Ctrl+D)
|
||||||
)
|
|
||||||
if not line:
|
|
||||||
break
|
break
|
||||||
|
|
||||||
self.current_input = line.rstrip('\n')
|
self.current_input = line.rstrip('\n')
|
||||||
|
|||||||
Reference in new issue
Block a user