407 lines
14 KiB
Python
407 lines
14 KiB
Python
"""
|
|
元宝全自动代理 - Solo 监听模式
|
|
|
|
使用方式：
|
|
1. 启动 Agent
|
|
2. 在元宝网页中直接对话
|
|
3. Agent 自动拦截 ACTION 标记并执行
|
|
4. 结果自动返回给元宝
|
|
|
|
无需终端交互，完全后台运行。
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import re
|
|
import subprocess
|
|
import time
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Optional, Any
|
|
from enum import Enum, auto
|
|
|
|
from playwright.async_api import async_playwright, Page, BrowserContext
|
|
|
|
YUANBAO_URL = "https://yuanbao.tencent.com/chat"
|
|
|
|
|
|
SOLO_SYSTEM_PROMPT = """你是一个自主任务执行助手。
|
|
|
|
## 可用操作
|
|
- `command`: 执行 shell 命令
|
|
- `powershell`: 执行 PowerShell
|
|
- `read_file`: 读取文件，cmd为路径
|
|
- `write_file`: 写入文件，cmd格式为 "路径::内容"
|
|
- `list_dir`: 列出目录
|
|
- `search`: 搜索文件内容
|
|
- `complete`: 任务完成
|
|
|
|
## 输出格式（严格按此格式）
|
|
<<<ACTION
|
|
{"type": "命令类型", "cmd": "具体命令", "explanation": "解释"}
|
|
ACTION>>>
|
|
|
|
任务完成时：
|
|
<<<ACTION
|
|
{"type": "complete", "explanation": "总结"}
|
|
ACTION>>>
|
|
|
|
## 规则
|
|
1. 每次只输出一个 ACTION 块
|
|
2. 等待系统返回结果后再继续
|
|
3. 不要输出 ACTION 块以外的操作指令
|
|
"""
|
|
|
|
|
|
class ExecutionEngine:
|
|
async def execute(self, action_type: str, command: str, working_dir: str = "") -> dict:
|
|
cwd = Path(working_dir).expanduser() if working_dir else Path.home()
|
|
|
|
handlers = {
|
|
"command": self._run_shell,
|
|
"powershell": self._run_powershell,
|
|
"read_file": self._read_file,
|
|
"write_file": self._write_file,
|
|
"list_dir": self._list_dir,
|
|
"search": self._search,
|
|
}
|
|
|
|
handler = handlers.get(action_type, self._run_shell)
|
|
|
|
try:
|
|
if asyncio.iscoroutinefunction(handler):
|
|
result = await handler(command, cwd)
|
|
else:
|
|
result = handler(command, cwd)
|
|
return {"success": True, "output": result}
|
|
except Exception as e:
|
|
return {"success": False, "output": f"Error: {e}"}
|
|
|
|
async def _run_shell(self, cmd: str, cwd: Path) -> str:
|
|
if sys.platform == "win32":
|
|
proc = await asyncio.create_subprocess_shell(
|
|
cmd, stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE, cwd=cwd,
|
|
creationflags=subprocess.CREATE_NO_WINDOW
|
|
)
|
|
else:
|
|
proc = await asyncio.create_subprocess_shell(
|
|
cmd, stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE, cwd=cwd
|
|
)
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=60)
|
|
output = stdout.decode('utf-8', errors='replace')
|
|
if stderr:
|
|
output += "\n[stderr] " + stderr.decode('utf-8', errors='replace')
|
|
return output
|
|
|
|
async def _run_powershell(self, cmd: str, cwd: Path) -> str:
|
|
ps_cmd = f'pwsh -NoProfile -Command "{cmd.replace(chr(34), chr(92)+chr(34))}"'
|
|
return await self._run_shell(ps_cmd, cwd)
|
|
|
|
def _read_file(self, path: str, cwd: Path) -> str:
|
|
file_path = Path(path).expanduser()
|
|
if not file_path.is_absolute():
|
|
file_path = cwd / file_path
|
|
return file_path.read_text(encoding='utf-8', errors='replace')
|
|
|
|
def _write_file(self, path_content: str, cwd: Path) -> str:
|
|
if "::" in path_content:
|
|
path, content = path_content.split("::", 1)
|
|
file_path = Path(path).expanduser()
|
|
if not file_path.is_absolute():
|
|
file_path = cwd / file_path
|
|
file_path.write_text(content, encoding='utf-8')
|
|
return f"Written: {file_path}"
|
|
return "Format error, use: path::content"
|
|
|
|
def _list_dir(self, path: str, cwd: Path) -> str:
|
|
dir_path = Path(path).expanduser() if path else cwd
|
|
if not dir_path.is_absolute():
|
|
dir_path = cwd / dir_path
|
|
|
|
lines = [f"Dir: {dir_path}"]
|
|
try:
|
|
for item in sorted(dir_path.iterdir()):
|
|
prefix = "[D]" if item.is_dir() else "[F]"
|
|
size = ""
|
|
if item.is_file():
|
|
size_bytes = item.stat().st_size
|
|
size = f" ({self._format_size(size_bytes)})"
|
|
lines.append(f" {prefix} {item.name}{size}")
|
|
except Exception as e:
|
|
lines.append(f" [Error] {e}")
|
|
return "\n".join(lines)
|
|
|
|
def _format_size(self, size: int) -> str:
|
|
for unit in ['B', 'KB', 'MB', 'GB']:
|
|
if size < 1024:
|
|
return f"{size:.1f}{unit}"
|
|
size /= 1024
|
|
return f"{size:.1f}TB"
|
|
|
|
async def _search(self, pattern: str, cwd: Path) -> str:
|
|
results = []
|
|
try:
|
|
for file_path in cwd.rglob("*"):
|
|
if file_path.is_file() and file_path.stat().st_size < 10*1024*1024:
|
|
try:
|
|
content = file_path.read_text(encoding='utf-8', errors='ignore')
|
|
if pattern in content:
|
|
lines = [i+1 for i, line in enumerate(content.split("\n")) if pattern in line]
|
|
results.append(f"{file_path.relative_to(cwd)}: lines {lines[:3]}")
|
|
if len(results) >= 10:
|
|
break
|
|
except:
|
|
continue
|
|
except Exception as e:
|
|
return f"Search failed: {e}"
|
|
return "\n".join(results) if results else "No matches"
|
|
|
|
|
|
class YuanbaoController:
|
|
def __init__(self, page: Page):
|
|
self.page = page
|
|
self.last_message = ""
|
|
|
|
async def send(self, content: str):
|
|
selectors = [
|
|
"div[contenteditable='true']",
|
|
"div[contenteditable='true'][data-placeholder]",
|
|
"textarea",
|
|
"[role='textbox']",
|
|
]
|
|
input_box = None
|
|
for sel in selectors:
|
|
try:
|
|
input_box = await self.page.wait_for_selector(sel, timeout=2000)
|
|
if input_box:
|
|
break
|
|
except:
|
|
continue
|
|
|
|
if not input_box:
|
|
print("[Error] Input not found")
|
|
return False
|
|
|
|
await input_box.click()
|
|
await asyncio.sleep(0.2)
|
|
await self.page.keyboard.press("Control+a")
|
|
await self.page.keyboard.press("Delete")
|
|
await asyncio.sleep(0.2)
|
|
|
|
try:
|
|
await input_box.fill(content)
|
|
except Exception as e:
|
|
print(f"[Warning] fill failed: {e}, using type")
|
|
await input_box.type(content, delay=1)
|
|
|
|
await asyncio.sleep(0.3)
|
|
|
|
try:
|
|
send_btn = self.page.locator("button[class*='send']").last
|
|
if await send_btn.count() > 0 and await send_btn.is_visible():
|
|
await send_btn.click()
|
|
else:
|
|
await input_box.press("Enter")
|
|
except:
|
|
await input_box.press("Enter")
|
|
|
|
return True
|
|
|
|
async def get_latest_ai_message(self) -> Optional[str]:
|
|
"""获取最新 AI 消息"""
|
|
try:
|
|
text = await self.page.evaluate("""
|
|
() => {
|
|
const selectors = [
|
|
'.agent-chat__list__item--ai',
|
|
'[class*="agent-chat"] [class*="ai"]',
|
|
'[class*="message"][class*="ai"]',
|
|
'[class*="assistant"]',
|
|
'.chat-item-ai',
|
|
];
|
|
|
|
for (const sel of selectors) {
|
|
const items = document.querySelectorAll(sel);
|
|
if (items.length > 0) {
|
|
const last = items[items.length - 1];
|
|
const textSelectors = [
|
|
'.agent-chat__speech-text',
|
|
'[class*="text"]',
|
|
'[class*="content"]',
|
|
'p',
|
|
'div',
|
|
];
|
|
for (const ts of textSelectors) {
|
|
const el = last.querySelector(ts);
|
|
if (el && el.innerText) {
|
|
return el.innerText.trim();
|
|
}
|
|
}
|
|
if (last.innerText) {
|
|
return last.innerText.trim();
|
|
}
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
""")
|
|
return text
|
|
except Exception as e:
|
|
print(f"[Error] get_latest_ai_message: {e}")
|
|
return None
|
|
|
|
|
|
class SoloAgent:
|
|
def __init__(self):
|
|
self.page: Optional[Page] = None
|
|
self.controller: Optional[YuanbaoController] = None
|
|
self.engine = ExecutionEngine()
|
|
self.last_ai_msg = ""
|
|
self.is_running = False
|
|
|
|
async def run(self):
|
|
print("=" * 60)
|
|
print(" Yuanbao Solo Agent - 监听模式")
|
|
print("=" * 60)
|
|
print("\n使用方式:")
|
|
print(" 1. 在元宝网页中直接对话")
|
|
print(" 2. Agent 自动拦截并执行")
|
|
print(" 3. 结果自动返回")
|
|
print("\n按 Ctrl+C 退出\n")
|
|
|
|
async with async_playwright() as p:
|
|
user_data = str(Path.home() / ".yuanbao_solo_profile")
|
|
context = await p.chromium.launch_persistent_context(
|
|
user_data, channel="chrome", headless=False,
|
|
args=["--start-maximized"], no_viewport=True
|
|
)
|
|
|
|
self.page = context.pages[0] if context.pages else await context.new_page()
|
|
await self.page.goto(YUANBAO_URL)
|
|
|
|
print(f"[Agent] 已打开元宝: {YUANBAO_URL}")
|
|
print("[Agent] 请登录并创建新对话\n")
|
|
|
|
input("登录完成后按回车继续...")
|
|
|
|
self.controller = YuanbaoController(self.page)
|
|
|
|
# 发送系统提示词
|
|
print("[Agent] 发送系统提示词...")
|
|
await self.controller.send(f"[系统初始化]\n\n{SOLO_SYSTEM_PROMPT}")
|
|
print("[Agent] 系统提示词已发送")
|
|
print("[Agent] 开始监听...\n")
|
|
|
|
# 初始化 last_ai_msg
|
|
await asyncio.sleep(3)
|
|
self.last_ai_msg = await self.controller.get_latest_ai_message() or ""
|
|
|
|
# 开始监听循环
|
|
self.is_running = True
|
|
await self._listen_loop()
|
|
|
|
async def _listen_loop(self):
|
|
"""持续监听元宝输出"""
|
|
while self.is_running:
|
|
try:
|
|
# 获取当前最新消息
|
|
current_msg = await self.controller.get_latest_ai_message()
|
|
|
|
if current_msg and current_msg != self.last_ai_msg:
|
|
print(f"\n[检测到新消息] {current_msg[:200]}...")
|
|
self.last_ai_msg = current_msg
|
|
|
|
# 检查是否包含 ACTION
|
|
action = self._parse_action(current_msg)
|
|
|
|
if action:
|
|
print(f"[解析成功] 动作: {action.get('type')}")
|
|
|
|
if action.get('type') == 'complete':
|
|
print("[任务完成]")
|
|
continue
|
|
|
|
# 执行动作
|
|
print(f"[执行] {action.get('type')}: {action.get('cmd', '')}")
|
|
result = await self.engine.execute(
|
|
action.get('type', 'command'),
|
|
action.get('cmd', ''),
|
|
""
|
|
)
|
|
|
|
print(f"[执行结果] Success: {result['success']}")
|
|
|
|
# 发送结果回元宝
|
|
result_msg = f"""[系统执行结果]
|
|
操作: {action.get('type')}
|
|
命令: {action.get('cmd', '')}
|
|
成功: {result['success']}
|
|
|
|
输出:
|
|
```
|
|
{result['output'][:1500]}
|
|
```
|
|
|
|
请继续下一步或完成任务。"""
|
|
|
|
await asyncio.sleep(1)
|
|
await self.controller.send(result_msg)
|
|
print("[已发送结果到元宝]")
|
|
else:
|
|
print("[无ACTION标记，忽略]")
|
|
|
|
await asyncio.sleep(1)
|
|
|
|
except Exception as e:
|
|
print(f"[错误] {e}")
|
|
await asyncio.sleep(2)
|
|
|
|
def _parse_action(self, text: str) -> Optional[dict]:
|
|
"""解析 ACTION 块"""
|
|
if not text:
|
|
return None
|
|
|
|
# 调试输出
|
|
if "<<<ACTION" in text:
|
|
print(f"[Debug] Found <<<ACTION marker")
|
|
|
|
# 支持多行格式，包括空行
|
|
# 格式1: <<<ACTION\n{...}\nACTION>>>
|
|
# 格式2: <<<ACTION{...}ACTION>>>
|
|
pattern = r'<<<ACTION\s*\n?(\{[\s\S]*?\})\s*ACTION>>>'
|
|
match = re.search(pattern, text, re.DOTALL)
|
|
if match:
|
|
try:
|
|
json_str = match.group(1).strip()
|
|
print(f"[Debug] Extracted JSON: {json_str[:100]}...")
|
|
return json.loads(json_str)
|
|
except json.JSONDecodeError as e:
|
|
print(f"[Debug] JSON parse error: {e}")
|
|
pass
|
|
|
|
# 备用：更宽松的匹配
|
|
if "<<<ACTION" in text and "ACTION>>>" in text:
|
|
try:
|
|
start = text.find("<<<ACTION") + len("<<<ACTION")
|
|
end = text.find("ACTION>>>")
|
|
content = text[start:end].strip()
|
|
if content.startswith("{") and content.endswith("}"):
|
|
print(f"[Debug] Fallback extraction: {content[:100]}...")
|
|
return json.loads(content)
|
|
except Exception as e:
|
|
print(f"[Debug] Fallback error: {e}")
|
|
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
agent = SoloAgent()
|
|
try:
|
|
asyncio.run(agent.run())
|
|
except KeyboardInterrupt:
|
|
print("\n[Agent] 已停止")
|