Files
blackbird/aibox/client.py
T
2026-07-27 21:39:53 +08:00

667 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
AI智能助手 - 命令行前端
"""
import asyncio
import json
import sys
import os
import threading
import hashlib
from datetime import datetime
from typing import Optional, Dict, Any, List
import websockets
from colorama import init, Fore, Style as ColoramaStyle
init(autoreset=True)
class Colors:
USER = Fore.CYAN
AI = Fore.GREEN
TASK = Fore.YELLOW
SYSTEM = Fore.BLUE
ERROR = Fore.RED
INFO = Fore.MAGENTA
RESET = ColoramaStyle.RESET_ALL
BOLD = ColoramaStyle.BRIGHT
GREEN = Fore.GREEN
RED = Fore.RED
class WebSocketClient:
def __init__(self, uri: str, callback):
self.uri = uri
self.websocket = None
self.connected = False
self.callback = callback
async def connect(self):
try:
self.websocket = await websockets.connect(self.uri)
self.connected = True
return True
except Exception as e:
self.callback("error", f"连接失败: {e}")
return False
async def disconnect(self):
if self.websocket:
await self.websocket.close()
self.connected = False
async def send_message(self, message: str):
if not self.websocket or not self.connected:
return False
try:
if message.startswith('/'):
data = {'type': 'command', 'command': message}
else:
data = {'type': 'message', 'content': message}
await self.websocket.send(json.dumps(data))
return True
except Exception as e:
self.callback("error", f"发送失败: {e}")
return False
async def receive_messages(self):
try:
async for message in self.websocket:
try:
data = json.loads(message)
await self._handle_message(data)
except json.JSONDecodeError:
self.callback("chunk", message)
except websockets.exceptions.ConnectionClosed:
self.connected = False
self.callback("system", "❌ 连接已断开")
except Exception as e:
self.callback("error", f"接收消息错误: {e}")
async def _handle_message(self, data: dict):
msg_type = data.get('type', '')
content = data.get('content', '')
if msg_type == 'chunk':
self.callback("chunk", content)
elif msg_type == 'complete':
self.callback("complete", content)
elif msg_type == 'dialogue_ended':
self.callback("dialogue_ended", content)
elif msg_type == 'task_process':
msg_subtype = data.get('msg_type', '')
task_content = data.get('content', '')
task_title = data.get('task_title', '未知任务')
if msg_subtype == 'chunk':
self.callback("task_chunk", task_content, task_title)
elif msg_subtype == 'complete':
self.callback("task_complete", task_title)
elif msg_type == 'task_status':
task_id = data.get('task_id', 'unknown')
title = data.get('title', '未知')
status = data.get('status', 'unknown')
message = data.get('message', content)
status_icon = {'executing': '⚙️', 'complete': '✅', 'failed': '❌', 'added': '➕', 'batch_start': '🚀', 'batch_complete': '🏁'}.get(status, '📋')
self.callback("info", f"{status_icon} [任务 {task_id}] {title}: {message}")
elif msg_type == 'task_error':
self.callback("error", f"[任务 {data.get('task_id', 'unknown')}] {data.get('task_title', '未知')}: {data.get('error', '未知错误')}")
elif msg_type == 'interrupt_ack':
self.callback("info", f"🛑 {content}")
elif msg_type == 'error':
self.callback("error", content)
elif msg_type == 'info':
self.callback("info", content)
elif msg_type == 'command_result':
self.callback("system", content)
elif msg_type == 'connection_ack':
self.callback("system", f"✅ {content}")
elif msg_type == 'notification':
self.callback("info", f"🔔 {content}")
else:
if content:
self.callback("chunk", content)
class SimpleCLI:
def __init__(self):
self.client: Optional[WebSocketClient] = None
self.connected = False
self.running = True
# AI消息缓冲
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
# 已显示消息的去重
self.displayed_messages: set = set()
# 锁和状态
self.lock = threading.Lock()
self.current_input = ""
# 流式输出状态
self.streaming_line = False # 是否正在流式输出行
self.in_tool_phase = False # 是否在工具调用阶段
self.waiting_for_response = False # 是否在等待AI响应
self.prompt_shown = True # 提示符是否已显示
# 命令列表
self.commands = {
'/help': '显示帮助信息',
'/connect': '连接到服务器',
'/disconnect': '断开连接',
'/quit': '退出程序',
'/clear': '清屏',
'/status': '显示连接状态',
'/interrupt': '发送中断信号(或按 Ctrl+C)',
}
def _get_prompt(self) -> str:
status = "● 已连接" if self.connected else "○ 未连接"
color = Colors.GREEN if self.connected else Colors.RED
return f"{color}{status}{Colors.RESET} >>> "
def _print_line(self, line: str = "", end: str = "\n"):
"""打印一行,并恢复提示符"""
with self.lock:
current = self.current_input
sys.stdout.write('\r' + ' ' * 120 + '\r')
sys.stdout.write(line + end)
sys.stdout.flush()
if not self.waiting_for_response and not self.in_tool_phase:
sys.stdout.write(self._get_prompt() + current)
sys.stdout.flush()
self.prompt_shown = True
self.streaming_line = False
def _print_prompt(self, clear_input: bool = False):
"""显示提示符"""
with self.lock:
if clear_input:
self.current_input = ""
self.waiting_for_response = False
self.in_tool_phase = False
self.prompt_shown = True
sys.stdout.write('\r' + ' ' * 120 + '\r')
sys.stdout.write(self._get_prompt() + self.current_input)
sys.stdout.flush()
self.streaming_line = False
def _print_stream(self, text: str):
"""流式打印文本(不换行,不显示提示符)"""
with self.lock:
# 保存当前输入
current = self.current_input
if not self.streaming_line:
# 首次流式输出,先换行再打印前缀
sys.stdout.write('\n')
sys.stdout.write(f"{Colors.AI}🤖 AI助手: ")
self.streaming_line = True
self.prompt_shown = False
# 打印文本
sys.stdout.write(text)
sys.stdout.flush()
# 注意:不恢复提示符,保持在同一行
def _finish_stream(self):
"""完成流式输出,换行并恢复提示符"""
with self.lock:
if self.streaming_line:
sys.stdout.write('\n')
self.streaming_line = False
if not self.in_tool_phase and not self.waiting_for_response:
current = self.current_input
sys.stdout.write(self._get_prompt() + current)
sys.stdout.flush()
self.prompt_shown = True
def _clean_content(self, content: str) -> str:
if not content:
return ""
if '💬 助手:' in content:
content = content.split('💬 助手:', 1)[1].strip()
content = content.lstrip()
return content
def _show_streaming(self, content: str):
"""流式显示增量内容"""
if not content:
return
# 清理内容
clean = self._clean_content(content)
if not clean:
return
key = clean[:200]
if key in self.displayed_messages:
return
# 检测增量
if self.ai_previous and clean.startswith(self.ai_previous):
new_part = clean[len(self.ai_previous):]
if new_part:
# 如果是首次流式输出,显示前缀
if not self.ai_prefix_shown:
self._print_stream("")
self.ai_prefix_shown = True
self._print_stream(new_part)
self.ai_previous = clean
self.ai_full = clean
return
if self.ai_previous and self.ai_previous in clean:
idx = clean.find(self.ai_previous) + len(self.ai_previous)
new_part = clean[idx:]
if new_part:
if not self.ai_prefix_shown:
self._print_stream("")
self.ai_prefix_shown = True
self._print_stream(new_part)
self.ai_previous = clean
self.ai_full = clean
return
# 全新内容 - 显示前缀和内容
if not self.ai_prefix_shown:
self._print_stream("")
self.ai_prefix_shown = True
self._print_stream(clean)
self.ai_previous = clean
self.ai_full = clean
def display(self, msg_type: str, content: str, extra: Any = None):
"""显示消息"""
if msg_type == 'chunk':
self._handle_chunk(content)
elif msg_type == 'complete':
self._handle_complete(content)
elif msg_type == 'task_chunk':
self._handle_task_chunk(content, extra)
elif msg_type == 'task_complete':
self._handle_task_complete(extra)
elif msg_type == 'dialogue_ended':
with self.lock:
if self.streaming_line:
sys.stdout.write('\n')
self.streaming_line = False
self.in_tool_phase = False
self.waiting_for_response = False
self.prompt_shown = True
self._print_line(f"{Colors.SYSTEM}✅ {content}{Colors.RESET}")
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
self.displayed_messages.clear()
self.streaming_line = False
elif msg_type == 'system':
self._finish_stream()
self._print_line(f"{Colors.SYSTEM}{content}{Colors.RESET}")
elif msg_type == 'info':
self._finish_stream()
self._print_line(f"{Colors.INFO}ℹ️ {content}{Colors.RESET}")
elif msg_type == 'error':
self._finish_stream()
self._print_line(f"{Colors.ERROR}❌ {content}{Colors.RESET}")
def _handle_chunk(self, content: str):
"""处理流式块"""
# 检查是否是工具调用相关的消息
if content.startswith('🔧 调用工具'):
with self.lock:
self.in_tool_phase = True
self.prompt_shown = False
# 打印工具调用信息
if self.streaming_line:
sys.stdout.write('\n')
self.streaming_line = False
sys.stdout.write(f"{Colors.TASK}{content}{Colors.RESET}\n")
sys.stdout.flush()
return
# 检查是否是工具结果
if content.startswith(('✅ 对话记忆已保存', '💾 保存记忆', '❌', '⚠️')):
with self.lock:
# 打印工具结果
sys.stdout.write(f"{Colors.TASK}{content}{Colors.RESET}\n")
sys.stdout.flush()
return
# 如果之前在工具阶段,现在退出工具阶段
if self.in_tool_phase:
with self.lock:
self.in_tool_phase = False
# 工具结束后不立即显示提示符,等待AI继续响应
self.prompt_shown = False
# 系统消息
if content.startswith(('📩', '[', '📁', '🔧', '📝', '✅', 'ℹ️', '⏰')):
self._finish_stream()
self._print_line(f"{Colors.INFO}{content}{Colors.RESET}")
return
# 流式显示
self._show_streaming(content)
def _handle_complete(self, content: str):
"""处理完成事件 - 只负责结束流式输出,不重复显示内容"""
# 如果之前在工具阶段,现在退出
if self.in_tool_phase:
with self.lock:
self.in_tool_phase = False
self.waiting_for_response = False
self.prompt_shown = True
current = self.current_input
sys.stdout.write(self._get_prompt() + current)
sys.stdout.flush()
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
self.streaming_line = False
return
# 如果流式输出正在进行,结束它
if self.streaming_line:
self._finish_stream()
# 如果内容不为空,检查是否已经通过流式显示了
if content and content.strip():
clean = self._clean_content(content)
if clean:
# 检查是否已通过流式显示
if self.ai_full and (clean in self.ai_full or self.ai_full in clean):
# 已经流式显示过了,不重复打印
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
self.streaming_line = False
# 确保提示符显示
with self.lock:
self.waiting_for_response = False
self.prompt_shown = True
current = self.current_input
sys.stdout.write(self._get_prompt() + current)
sys.stdout.flush()
return
# 检查是否已经在 displayed_messages 中
key = clean[:200]
if key in self.displayed_messages:
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
self.streaming_line = False
# 确保提示符显示
with self.lock:
self.waiting_for_response = False
self.prompt_shown = True
current = self.current_input
sys.stdout.write(self._get_prompt() + current)
sys.stdout.flush()
return
# 只有当内容没有被流式显示时才显示完整消息
self.displayed_messages.add(key)
# 如果之前没有流式输出,现在显示完整消息
self._print_line(f"{Colors.AI}🤖 AI助手: {clean}")
# 重置状态
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
self.streaming_line = False
self.waiting_for_response = False
# 确保提示符显示
with self.lock:
self.prompt_shown = True
current = self.current_input
sys.stdout.write('\r' + ' ' * 120 + '\r')
sys.stdout.write(self._get_prompt() + current)
sys.stdout.flush()
def _handle_task_chunk(self, content: str, task_title: str):
"""处理任务流式块"""
if not hasattr(self, '_task_buffers'):
self._task_buffers = {}
if task_title not in self._task_buffers:
self._task_buffers[task_title] = ""
self._finish_stream()
self._print_line(f"{Colors.TASK}⚡ 任务 - {task_title}: ", end="")
self.streaming_line = True
prev = self._task_buffers[task_title]
if content.startswith(prev):
new_part = content[len(prev):]
if new_part:
with self.lock:
sys.stdout.write(new_part)
sys.stdout.flush()
self._task_buffers[task_title] = content
elif prev in content:
idx = content.find(prev) + len(prev)
new_part = content[idx:]
if new_part:
with self.lock:
sys.stdout.write(new_part)
sys.stdout.flush()
self._task_buffers[task_title] = content
else:
with self.lock:
sys.stdout.write(content)
sys.stdout.flush()
self._task_buffers[task_title] = content
def _handle_task_complete(self, task_title: str):
"""处理任务完成"""
self._finish_stream()
if hasattr(self, '_task_buffers') and task_title in self._task_buffers:
content = self._task_buffers[task_title]
if content:
clean = content.strip()
key = clean[:200]
if key not in self.displayed_messages:
self.displayed_messages.add(key)
self._print_line(f"{Colors.TASK}⚡ 任务 - {task_title}: {clean}")
del self._task_buffers[task_title]
self.streaming_line = False
def print_banner(self):
banner = f"""
{Colors.BOLD}{Colors.AI}
╔══════════════════════════════════════════════════════════╗
║ AI智能助手 - 命令行前端 ║
║ WebSocket 直连 + CTRL+C 中断 ║
╚══════════════════════════════════════════════════════════╝
{Colors.RESET}
"""
self._print_line(banner)
self._print_line(f"{Colors.INFO}📌 输入 /help 查看可用命令")
self._print_line(f"{Colors.INFO}📌 按 Ctrl+C 发送中断信号 (而非退出)")
self._print_line(f"{Colors.INFO}📌 按 Ctrl+D 退出程序{Colors.RESET}")
self._print_line("")
self._print_line("=" * 60)
self._print_line("")
self._print_prompt(clear_input=True)
def show_help(self):
self._print_line("")
self._print_line(f"{Colors.BOLD}可用命令:{Colors.RESET}")
self._print_line("")
self._print_line(f"{Colors.USER}本地命令:{Colors.RESET}")
self._print_line(" /connect <url> 连接到WebSocket服务器")
self._print_line(" /disconnect 断开连接")
self._print_line(" /status 显示连接状态")
self._print_line(" /clear 清屏")
self._print_line(" /interrupt 发送中断信号 (或按 Ctrl+C)")
self._print_line(" /quit 退出程序 (或按 Ctrl+D)")
self._print_line(" /help 显示此帮助")
self._print_line("")
self._print_line(f"{Colors.AI}服务器命令 (需要先连接):{Colors.RESET}")
for cmd, desc in self.commands.items():
if cmd not in ['/help', '/connect', '/disconnect', '/quit', '/clear', '/status', '/interrupt']:
self._print_line(f" {cmd:<15} {desc}")
self._print_prompt()
async def handle_command(self, command: str):
parts = command.split()
cmd = parts[0].lower()
if cmd == '/help':
self.show_help()
elif cmd == '/connect':
if len(parts) < 2:
self.display("error", "用法: /connect ws://host:port")
return
uri = parts[1]
if self.client and self.client.connected:
await self.client.disconnect()
self.client = WebSocketClient(uri, self.display)
if await self.client.connect():
self.connected = True
self.display("system", f"✅ 已连接到 {uri}")
asyncio.create_task(self.client.receive_messages())
else:
self.client = None
self.connected = False
elif cmd == '/disconnect':
if self.client:
await self.client.disconnect()
self.client = None
self.connected = False
self.display("system", "已断开连接")
else:
self.display("error", "未连接到服务器")
elif cmd == '/status':
if self.connected:
self.display("system", f"✅ 已连接: {self.client.uri}")
else:
self.display("system", "未连接")
elif cmd == '/clear':
os.system('clear' if os.name == 'posix' else 'cls')
self.displayed_messages.clear()
self.current_input = ""
self.print_banner()
elif cmd == '/interrupt':
if self.client and self.client.connected:
await self._send_interrupt()
self.display("info", "🛑 已发送中断信号")
else:
self.display("error", "未连接到服务器")
elif cmd == '/quit':
self.display("system", "👋 再见!")
self.running = False
if self.client:
await self.client.disconnect()
return
elif cmd in self.commands:
if self.client and self.client.connected:
await self.client.send_message(command)
self.display("system", f"📤 发送命令: {command}")
else:
self.display("error", "未连接到服务器")
else:
self.display("error", f"未知命令: {cmd}")
self._print_prompt(clear_input=True)
async def _send_interrupt(self):
if self.client and self.client.connected:
interrupt_msg = {"type": "interrupt", "content": "用户按下了 Ctrl+C 中断当前操作"}
try:
await self.client.websocket.send(json.dumps(interrupt_msg))
except Exception as e:
self.display("error", f"发送中断失败: {e}")
async def send_user_message(self, message: str):
if not self.client or not self.client.connected:
self.display("error", "未连接到服务器")
return
# 重置流式状态
with self.lock:
if self.streaming_line:
sys.stdout.write('\n')
self.streaming_line = False
self.waiting_for_response = True
self.in_tool_phase = False
self.prompt_shown = False
# 【关键】清除当前输入,避免显示之前输入的内容
self.current_input = ""
self.ai_previous = ""
self.ai_full = ""
self.ai_prefix_shown = False
self.display("system", f"{Colors.USER}👤 用户: {message}{Colors.RESET}")
await self.client.send_message(message)
# 不在这里显示提示符,等待AI响应结束后再显示
async def run(self):
self.print_banner()
self._task_buffers = {}
while self.running:
try:
line = await asyncio.get_event_loop().run_in_executor(
None, lambda: sys.stdin.readline()
)
if not line:
break
self.current_input = line.rstrip('\n')
line = line.strip()
if not line:
self._print_prompt(clear_input=True)
continue
if line.startswith('/'):
await self.handle_command(line)
else:
await self.send_user_message(line)
except KeyboardInterrupt:
if self.client and self.client.connected:
await self._send_interrupt()
self.display("info", "🛑 已发送中断信号")
else:
self.display("error", "未连接到服务器")
self._print_prompt(clear_input=True)
continue
except Exception as e:
if self.running:
self.display("error", f"错误: {e}")
break
async def main():
default_uri = "ws://localhost:8765"
uri = sys.argv[1] if len(sys.argv) > 1 else default_uri
cli = SimpleCLI()
cli.client = WebSocketClient(uri, cli.display)
if await cli.client.connect():
cli.connected = True
cli.display("system", f"✅ 已连接到 {uri}")
asyncio.create_task(cli.client.receive_messages())
else:
cli.display("error", f"无法连接到 {uri}")
cli.display("info", f"你可以通过命令行参数指定地址: python client.py ws://your-server:8765")
try:
await cli.run()
except KeyboardInterrupt:
print(f"\n{Colors.SYSTEM}程序已退出{Colors.RESET}")
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
print(f"\n{Colors.SYSTEM}程序已退出{Colors.RESET}")
except Exception as e:
print(f"\n{Colors.ERROR}错误: {e}{Colors.RESET}")