374 lines
12 KiB
Python
374 lines
12 KiB
Python
"""
|
||
元宝全自动代理 MVP
|
||
验证思路：通过 Playwright 拦截元宝对话，识别本地调用意图，
|
||
执行本地操作后将结果注入回对话。
|
||
|
||
核心机制：
|
||
1. 监听元宝输入框，检测特殊指令标记（如 /local 或 @agent）
|
||
2. 拦截用户输入，在本地执行
|
||
3. 将执行结果作为"用户消息"注入回对话
|
||
4. 元宝基于结果继续回复
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Optional, Callable
|
||
|
||
from playwright.async_api import async_playwright, Page, BrowserContext
|
||
|
||
YUANBAO_URL = "https://yuanbao.tencent.com/chat"
|
||
|
||
# 本地工具注册表
|
||
TOOLS: dict[str, Callable] = {}
|
||
|
||
|
||
def register_tool(name: str):
|
||
"""工具装饰器"""
|
||
def decorator(func: Callable):
|
||
TOOLS[name] = func
|
||
return func
|
||
return decorator
|
||
|
||
|
||
# ============ 本地工具实现 ============
|
||
|
||
@register_tool("read_file")
|
||
def tool_read_file(path: str) -> str:
|
||
"""读取本地文件内容"""
|
||
try:
|
||
file_path = Path(path).expanduser().resolve()
|
||
if not file_path.exists():
|
||
return f"[错误] 文件不存在: {path}"
|
||
content = file_path.read_text(encoding="utf-8", errors="ignore")
|
||
# 截断过长内容
|
||
if len(content) > 5000:
|
||
content = content[:5000] + f"\n... (已截断，共 {len(content)} 字符)"
|
||
return f"[文件内容: {path}]\n```\n{content}\n```"
|
||
except Exception as e:
|
||
return f"[错误] 读取文件失败: {e}"
|
||
|
||
|
||
@register_tool("list_dir")
|
||
def tool_list_dir(path: str = ".") -> str:
|
||
"""列出目录内容"""
|
||
try:
|
||
dir_path = Path(path).expanduser().resolve()
|
||
if not dir_path.exists():
|
||
return f"[错误] 目录不存在: {path}"
|
||
items = []
|
||
for item in sorted(dir_path.iterdir()):
|
||
item_type = "📁" if item.is_dir() else "📄"
|
||
items.append(f"{item_type} {item.name}")
|
||
return f"[目录列表: {path}]\n" + "\n".join(items[:50])
|
||
except Exception as e:
|
||
return f"[错误] 列出目录失败: {e}"
|
||
|
||
|
||
@register_tool("run_command")
|
||
def tool_run_command(command: str) -> str:
|
||
"""执行本地命令"""
|
||
try:
|
||
result = subprocess.run(
|
||
command,
|
||
shell=True,
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=30,
|
||
encoding="utf-8",
|
||
errors="ignore"
|
||
)
|
||
output = result.stdout
|
||
if result.stderr:
|
||
output += f"\n[stderr]\n{result.stderr}"
|
||
# 截断过长输出
|
||
if len(output) > 5000:
|
||
output = output[:5000] + f"\n... (已截断，共 {len(output)} 字符)"
|
||
return f"[命令执行: {command}]\n返回码: {result.returncode}\n```\n{output}\n```"
|
||
except subprocess.TimeoutExpired:
|
||
return f"[错误] 命令执行超时: {command}"
|
||
except Exception as e:
|
||
return f"[错误] 执行命令失败: {e}"
|
||
|
||
|
||
@register_tool("web_search")
|
||
def tool_web_search(query: str) -> str:
|
||
"""模拟网络搜索（MVP 阶段返回提示）"""
|
||
return f"[搜索: {query}]\nMVP 阶段暂未接入真实搜索，请直接询问元宝。"
|
||
|
||
|
||
# ============ 指令解析器 ============
|
||
|
||
class CommandParser:
|
||
"""解析用户输入中的本地调用指令"""
|
||
|
||
# 指令模式：@agent:tool_name(args) 或 /local tool_name args
|
||
PATTERNS = [
|
||
# @agent:read_file("/path/to/file")
|
||
r'@agent:(\w+)\s*\(\s*["\']?([^"\')]*)["\']?\s*\)',
|
||
# /local read_file /path/to/file
|
||
r'/local\s+(\w+)\s+(.+?)(?:\n|$)',
|
||
# ```local
|
||
# read_file /path/to/file
|
||
# ```
|
||
r'```local\s*\n(\w+)\s+(.+?)\n```',
|
||
]
|
||
|
||
def parse(self, text: str) -> Optional[dict]:
|
||
"""解析指令，返回工具调用信息或 None"""
|
||
text = text.strip()
|
||
|
||
# 模式 1: @agent:tool_name(args)
|
||
match = re.search(self.PATTERNS[0], text)
|
||
if match:
|
||
return {
|
||
"tool": match.group(1),
|
||
"args": [match.group(2)],
|
||
"raw": match.group(0)
|
||
}
|
||
|
||
# 模式 2: /local tool_name args
|
||
match = re.search(self.PATTERNS[1], text)
|
||
if match:
|
||
return {
|
||
"tool": match.group(1),
|
||
"args": [match.group(2).strip()],
|
||
"raw": match.group(0)
|
||
}
|
||
|
||
# 模式 3: ```local ... ```
|
||
match = re.search(self.PATTERNS[2], text, re.DOTALL)
|
||
if match:
|
||
return {
|
||
"tool": match.group(1),
|
||
"args": [match.group(2).strip()],
|
||
"raw": match.group(0)
|
||
}
|
||
|
||
return None
|
||
|
||
|
||
# ============ 元宝页面操作 ============
|
||
|
||
class YuanbaoController:
|
||
"""控制元宝网页的输入输出"""
|
||
|
||
def __init__(self, page: Page):
|
||
self.page = page
|
||
self.parser = CommandParser()
|
||
|
||
async def find_input_box(self):
|
||
"""查找输入框"""
|
||
selectors = [
|
||
"div[contenteditable='true'][data-placeholder]",
|
||
"div[contenteditable='true']",
|
||
"textarea[placeholder]",
|
||
"textarea",
|
||
]
|
||
for sel in selectors:
|
||
try:
|
||
el = await self.page.wait_for_selector(sel, timeout=3000)
|
||
if el:
|
||
return el
|
||
except Exception:
|
||
continue
|
||
return None
|
||
|
||
async def inject_message(self, content: str):
|
||
"""将内容注入到输入框并发送"""
|
||
input_box = await self.find_input_box()
|
||
if not input_box:
|
||
print("[MVP] 未找到输入框")
|
||
return False
|
||
|
||
# 点击输入框
|
||
await input_box.click()
|
||
await asyncio.sleep(0.3)
|
||
|
||
# 清空现有内容
|
||
await self.page.keyboard.press("Control+a")
|
||
await self.page.keyboard.press("Delete")
|
||
await asyncio.sleep(0.2)
|
||
|
||
# 输入内容
|
||
# 分段输入避免过长
|
||
max_chunk = 500
|
||
for i in range(0, len(content), max_chunk):
|
||
chunk = content[i:i+max_chunk]
|
||
await input_box.type(chunk, delay=5)
|
||
|
||
await asyncio.sleep(0.5)
|
||
|
||
# 发送
|
||
await self._click_send(input_box)
|
||
return True
|
||
|
||
async def _click_send(self, input_box):
|
||
"""点击发送按钮"""
|
||
send_selectors = [
|
||
"button[class*='send']",
|
||
"button[aria-label*='发送']",
|
||
"button[title*='发送']",
|
||
"div[class*='send-btn']",
|
||
]
|
||
for sel in send_selectors:
|
||
try:
|
||
btn = self.page.locator(sel).last
|
||
if await btn.is_visible():
|
||
await btn.click()
|
||
return
|
||
except Exception:
|
||
continue
|
||
# 回退到 Enter
|
||
await input_box.press("Enter")
|
||
|
||
async def get_last_user_message(self) -> Optional[str]:
|
||
"""获取最后一条用户消息"""
|
||
try:
|
||
# 查找用户消息
|
||
messages = await self.page.evaluate("""
|
||
() => {
|
||
const items = document.querySelectorAll('.agent-chat__list__item--user, [class*="user"]');
|
||
const last = items[items.length - 1];
|
||
return last ? last.innerText?.trim() : null;
|
||
}
|
||
""")
|
||
return messages
|
||
except Exception as e:
|
||
print(f"[MVP] 获取消息失败: {e}")
|
||
return None
|
||
|
||
async def intercept_and_execute(self, user_input: str) -> Optional[str]:
|
||
"""
|
||
拦截用户输入，如果包含本地指令则执行并返回结果
|
||
返回: 要注入回对话的内容，或 None（无需拦截）
|
||
"""
|
||
cmd = self.parser.parse(user_input)
|
||
if not cmd:
|
||
return None
|
||
|
||
tool_name = cmd["tool"]
|
||
args = cmd["args"]
|
||
|
||
print(f"[MVP] 检测到本地指令: {tool_name}({args})")
|
||
|
||
# 执行工具
|
||
if tool_name not in TOOLS:
|
||
return f"[错误] 未知工具: {tool_name}。可用工具: {list(TOOLS.keys())}"
|
||
|
||
try:
|
||
result = TOOLS[tool_name](*args)
|
||
# 包装成自然对话格式
|
||
return f"[本地执行结果]\n{result}\n\n请基于以上结果继续回答。"
|
||
except Exception as e:
|
||
return f"[错误] 工具执行失败: {e}"
|
||
|
||
|
||
# ============ 主循环 ============
|
||
|
||
class YuanbaoAgentMVP:
|
||
"""元宝全自动代理 MVP"""
|
||
|
||
def __init__(self):
|
||
self.context: Optional[BrowserContext] = None
|
||
self.page: Optional[Page] = None
|
||
self.controller: Optional[YuanbaoController] = None
|
||
self._last_processed: Optional[str] = None
|
||
|
||
async def run(self):
|
||
"""启动 MVP"""
|
||
print("=" * 60)
|
||
print(" 元宝全自动代理 MVP")
|
||
print("=" * 60)
|
||
print("\n使用说明:")
|
||
print(" 1. 在元宝对话中输入特殊指令触发本地调用:")
|
||
print(" - @agent:read_file('/path/to/file')")
|
||
print(" - /local list_dir /path/to/dir")
|
||
print(" - ```local\nrun_command ls -la\n```")
|
||
print(" 2. 系统将拦截指令，执行本地操作")
|
||
print(" 3. 执行结果将注入回对话，元宝基于结果继续")
|
||
print("\n按 Ctrl+C 退出\n")
|
||
|
||
async with async_playwright() as p:
|
||
# 启动浏览器
|
||
print("[MVP] 启动浏览器...")
|
||
user_data_dir = str(Path.home() / ".yuanbao_agent_profile")
|
||
self.context = await p.chromium.launch_persistent_context(
|
||
user_data_dir=user_data_dir,
|
||
channel="chrome",
|
||
headless=False,
|
||
args=["--start-maximized"],
|
||
no_viewport=True,
|
||
)
|
||
|
||
self.page = self.context.pages[0] if self.context.pages else await self.context.new_page()
|
||
await self.page.goto(YUANBAO_URL)
|
||
|
||
print(f"[MVP] 已打开元宝: {YUANBAO_URL}")
|
||
print("[MVP] 请手动登录元宝，然后按回车继续...")
|
||
|
||
loop = asyncio.get_event_loop()
|
||
await loop.run_in_executor(None, input)
|
||
|
||
self.controller = YuanbaoController(self.page)
|
||
|
||
print("[MVP] 开始监听对话...")
|
||
await self._monitor_loop()
|
||
|
||
async def _monitor_loop(self):
|
||
"""监听循环"""
|
||
while True:
|
||
try:
|
||
await asyncio.sleep(2) # 每 2 秒检查一次
|
||
|
||
# 获取最后一条用户消息
|
||
last_msg = await self.controller.get_last_user_message()
|
||
if not last_msg:
|
||
continue
|
||
|
||
# 避免重复处理
|
||
if last_msg == self._last_processed:
|
||
continue
|
||
|
||
# 检查是否包含本地指令
|
||
result = await self.controller.intercept_and_execute(last_msg)
|
||
if result:
|
||
print(f"[MVP] 准备注入执行结果...")
|
||
# 等待元宝开始回复（可选）
|
||
await asyncio.sleep(1)
|
||
|
||
# 注入结果
|
||
success = await self.controller.inject_message(result)
|
||
if success:
|
||
print(f"[MVP] 结果已注入对话")
|
||
self._last_processed = last_msg
|
||
else:
|
||
print(f"[MVP] 注入失败")
|
||
else:
|
||
# 无指令，只是记录
|
||
self._last_processed = last_msg
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n[MVP] 用户中断")
|
||
break
|
||
except Exception as e:
|
||
print(f"[MVP] 监听错误: {e}")
|
||
await asyncio.sleep(5)
|
||
|
||
|
||
async def main():
|
||
agent = YuanbaoAgentMVP()
|
||
try:
|
||
await agent.run()
|
||
except Exception as e:
|
||
print(f"[MVP] 错误: {e}")
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|