523 lines
17 KiB
Python
523 lines
17 KiB
Python
"""
|
||
元宝全自动代理 v2 - 智能工具调用
|
||
|
||
核心改进：
|
||
1. 元宝主动输出工具调用标记（JSON 格式）
|
||
2. Agent 拦截标记，执行工具，将结果返回给元宝
|
||
3. 支持多轮工具调用直到任务完成
|
||
4. 更自然的对话体验
|
||
|
||
工作流程：
|
||
用户提问 → 元宝分析 → 输出工具调用标记 → Agent 拦截执行 →
|
||
结果返回元宝 → 元宝继续分析/输出新标记 → ... → 最终答案
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Optional, Callable, Any
|
||
|
||
from playwright.async_api import async_playwright, Page, BrowserContext
|
||
|
||
YUANBAO_URL = "https://yuanbao.tencent.com/chat"
|
||
|
||
# 工具调用标记格式
|
||
TOOL_CALL_START = "<<<TOOL_CALL"
|
||
TOOL_CALL_END = "TOOL_CALL>>>"
|
||
|
||
|
||
# ============ 数据模型 ============
|
||
|
||
@dataclass
|
||
class ToolCall:
|
||
"""工具调用请求"""
|
||
tool: str
|
||
args: list
|
||
call_id: str = ""
|
||
|
||
|
||
@dataclass
|
||
class ToolResult:
|
||
"""工具执行结果"""
|
||
call_id: str
|
||
success: bool
|
||
content: str
|
||
|
||
|
||
# ============ 工具注册表 ============
|
||
|
||
class ToolRegistry:
|
||
"""工具注册与管理"""
|
||
|
||
def __init__(self):
|
||
self._tools: dict[str, Callable] = {}
|
||
self._schemas: dict[str, dict] = {}
|
||
|
||
def register(self, name: str, schema: dict):
|
||
"""注册工具装饰器"""
|
||
def decorator(func: Callable):
|
||
self._tools[name] = func
|
||
self._schemas[name] = schema
|
||
return func
|
||
return decorator
|
||
|
||
def execute(self, tool_name: str, args: list) -> ToolResult:
|
||
"""执行工具"""
|
||
if tool_name not in self._tools:
|
||
return ToolResult(
|
||
call_id="",
|
||
success=False,
|
||
content=f"未知工具: {tool_name}。可用工具: {list(self._tools.keys())}"
|
||
)
|
||
|
||
try:
|
||
result = self._tools[tool_name](*args)
|
||
return ToolResult(call_id="", success=True, content=str(result))
|
||
except Exception as e:
|
||
return ToolResult(call_id="", success=False, content=f"执行失败: {e}")
|
||
|
||
def get_schemas(self) -> dict:
|
||
"""获取所有工具 schema"""
|
||
return self._schemas
|
||
|
||
|
||
# 全局工具注册表
|
||
registry = ToolRegistry()
|
||
|
||
|
||
# ============ 工具实现 ============
|
||
|
||
@registry.register("read_file", {
|
||
"description": "读取本地文件内容",
|
||
"parameters": [
|
||
{"name": "path", "type": "string", "description": "文件路径"}
|
||
]
|
||
})
|
||
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) > 8000:
|
||
content = content[:8000] + f"\n... (已截断，共 {len(content)} 字符)"
|
||
return content
|
||
except Exception as e:
|
||
return f"[错误] {e}"
|
||
|
||
|
||
@registry.register("write_file", {
|
||
"description": "写入内容到本地文件",
|
||
"parameters": [
|
||
{"name": "path", "type": "string", "description": "文件路径"},
|
||
{"name": "content", "type": "string", "description": "文件内容"}
|
||
]
|
||
})
|
||
def tool_write_file(path: str, content: str) -> str:
|
||
try:
|
||
file_path = Path(path).expanduser().resolve()
|
||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||
file_path.write_text(content, encoding="utf-8")
|
||
return f"[成功] 文件已写入: {path}"
|
||
except Exception as e:
|
||
return f"[错误] {e}"
|
||
|
||
|
||
@registry.register("list_dir", {
|
||
"description": "列出目录内容",
|
||
"parameters": [
|
||
{"name": "path", "type": "string", "description": "目录路径", "default": "."}
|
||
]
|
||
})
|
||
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()):
|
||
prefix = "[DIR]" if item.is_dir() else "[FILE]"
|
||
size = ""
|
||
if item.is_file():
|
||
try:
|
||
size_bytes = item.stat().st_size
|
||
size = f" ({size_bytes:,} bytes)"
|
||
except:
|
||
pass
|
||
items.append(f"{prefix} {item.name}{size}")
|
||
return "\n".join(items[:100])
|
||
except Exception as e:
|
||
return f"[错误] {e}"
|
||
|
||
|
||
@registry.register("run_command", {
|
||
"description": "执行本地命令",
|
||
"parameters": [
|
||
{"name": "command", "type": "string", "description": "要执行的命令"}
|
||
]
|
||
})
|
||
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] {result.stderr}"
|
||
if len(output) > 8000:
|
||
output = output[:8000] + f"\n... (已截断)"
|
||
return f"[退出码: {result.returncode}]\n{output}"
|
||
except subprocess.TimeoutExpired:
|
||
return "[错误] 命令执行超时"
|
||
except Exception as e:
|
||
return f"[错误] {e}"
|
||
|
||
|
||
@registry.register("search_files", {
|
||
"description": "搜索文件内容",
|
||
"parameters": [
|
||
{"name": "pattern", "type": "string", "description": "搜索模式"},
|
||
{"name": "path", "type": "string", "description": "搜索路径", "default": "."}
|
||
]
|
||
})
|
||
def tool_search_files(pattern: str, path: str = ".") -> str:
|
||
"""简单文件内容搜索"""
|
||
try:
|
||
import fnmatch
|
||
search_path = Path(path).expanduser().resolve()
|
||
results = []
|
||
|
||
for file_path in search_path.rglob("*"):
|
||
if file_path.is_file() and not any(p.startswith(".") for p in file_path.parts):
|
||
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(search_path)}: 行 {lines[:5]}")
|
||
if len(results) >= 20:
|
||
results.append("... (结果过多，已截断)")
|
||
break
|
||
except:
|
||
continue
|
||
|
||
return "\n".join(results) if results else "未找到匹配内容"
|
||
except Exception as e:
|
||
return f"[错误] {e}"
|
||
|
||
|
||
# ============ 系统提示词 ============
|
||
|
||
SYSTEM_PROMPT = f"""你是一个具有本地操作能力的 AI 助手。你可以调用本地工具来帮助用户完成任务。
|
||
|
||
可用工具及其 schema：
|
||
{json.dumps(registry.get_schemas(), indent=2, ensure_ascii=False)}
|
||
|
||
工具调用规则：
|
||
1. 当需要调用工具时，输出以下格式的标记：
|
||
<<<TOOL_CALL
|
||
{{"tool": "工具名", "args": ["参数1", "参数2", ...]}}
|
||
TOOL_CALL>>>
|
||
|
||
2. 每次只能调用一个工具
|
||
3. 等待工具执行结果后再继续
|
||
4. 最终结果以自然语言回复用户
|
||
|
||
示例对话：
|
||
用户：帮我看看桌面有什么文件
|
||
你：我来查看一下桌面文件。
|
||
<<<TOOL_CALL
|
||
{{"tool": "list_dir", "args": ["~/Desktop"]}}
|
||
TOOL_CALL>>>
|
||
[系统注入工具结果...]
|
||
你：桌面有以下文件：...
|
||
"""
|
||
|
||
|
||
# ============ 元宝页面控制器 ============
|
||
|
||
class YuanbaoController:
|
||
"""控制元宝网页交互"""
|
||
|
||
def __init__(self, page: Page):
|
||
self.page = page
|
||
self._last_ai_message: Optional[str] = None
|
||
|
||
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 send_message(self, content: str):
|
||
"""发送消息到元宝"""
|
||
input_box = await self.find_input_box()
|
||
if not input_box:
|
||
print("[v2] 未找到输入框")
|
||
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)
|
||
|
||
# 使用 fill 方法一次性填充内容（比分段 type 更可靠）
|
||
try:
|
||
await input_box.fill(content)
|
||
await asyncio.sleep(0.3)
|
||
except Exception as e:
|
||
# 兜底：逐字输入
|
||
print(f"[v2] fill 失败，使用 type: {e}")
|
||
await input_box.type(content, delay=1)
|
||
|
||
await asyncio.sleep(0.3)
|
||
|
||
# 发送
|
||
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*='发送']",
|
||
]
|
||
for sel in send_selectors:
|
||
try:
|
||
btn = self.page.locator(sel).last
|
||
if await btn.is_visible():
|
||
await btn.click()
|
||
return
|
||
except Exception:
|
||
continue
|
||
await input_box.press("Enter")
|
||
|
||
async def get_latest_ai_message(self) -> Optional[str]:
|
||
"""获取最新 AI 回复"""
|
||
try:
|
||
text = await self.page.evaluate("""
|
||
() => {
|
||
const items = document.querySelectorAll('.agent-chat__list__item--ai');
|
||
if (items.length === 0) return null;
|
||
|
||
// 找最后一个有 speech-text-box 的
|
||
let lastItem = null;
|
||
for (let i = items.length - 1; i >= 0; i--) {
|
||
if (items[i].querySelector('.agent-chat__speech-text-box')) {
|
||
lastItem = items[i];
|
||
break;
|
||
}
|
||
}
|
||
if (!lastItem) lastItem = items[items.length - 1];
|
||
|
||
const textBox = lastItem.querySelector('.agent-chat__speech-text');
|
||
return textBox ? textBox.innerText?.trim() : lastItem.innerText?.trim();
|
||
}
|
||
""")
|
||
return text
|
||
except Exception as e:
|
||
print(f"[v2] 获取 AI 消息失败: {e}")
|
||
return None
|
||
|
||
async def wait_for_new_ai_message(self, previous: Optional[str] = None, timeout: int = 120) -> Optional[str]:
|
||
"""等待新的 AI 消息"""
|
||
deadline = time.time() + timeout
|
||
|
||
while time.time() < deadline:
|
||
current = await self.get_latest_ai_message()
|
||
if current and current != previous:
|
||
return current
|
||
await asyncio.sleep(1)
|
||
|
||
return None
|
||
|
||
def parse_tool_call(self, text: str) -> Optional[ToolCall]:
|
||
"""解析工具调用标记"""
|
||
# 支持多行格式，允许标记和 JSON 之间有换行
|
||
pattern = rf'{re.escape(TOOL_CALL_START)}\s*(\{{.*?\}})\s*{re.escape(TOOL_CALL_END)}'
|
||
match = re.search(pattern, text, re.DOTALL)
|
||
|
||
if match:
|
||
try:
|
||
json_str = match.group(1).strip()
|
||
data = json.loads(json_str)
|
||
return ToolCall(
|
||
tool=data.get("tool", ""),
|
||
args=data.get("args", []),
|
||
call_id=data.get("call_id", "")
|
||
)
|
||
except json.JSONDecodeError as e:
|
||
print(f"[v2] JSON 解析失败: {e}, 内容: {match.group(1)[:200]}")
|
||
return None
|
||
return None
|
||
|
||
|
||
# ============ 主代理 ============
|
||
|
||
class YuanbaoAgent:
|
||
"""元宝全自动代理"""
|
||
|
||
def __init__(self):
|
||
self.context: Optional[BrowserContext] = None
|
||
self.page: Optional[Page] = None
|
||
self.controller: Optional[YuanbaoController] = None
|
||
self.conversation_history: list[dict] = []
|
||
|
||
async def run(self):
|
||
"""启动代理"""
|
||
print("=" * 60)
|
||
print(" 元宝全自动代理 v2")
|
||
print("=" * 60)
|
||
print("\n工作流程:")
|
||
print(" 1. 用户在元宝中提问")
|
||
print(" 2. 元宝输出工具调用标记")
|
||
print(" 3. Agent 拦截并执行本地操作")
|
||
print(" 4. 结果返回元宝，继续对话")
|
||
print("\n按 Ctrl+C 退出\n")
|
||
|
||
async with async_playwright() as p:
|
||
print("[Agent] 启动浏览器...")
|
||
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"[Agent] 已打开元宝: {YUANBAO_URL}")
|
||
print("[Agent] 请手动登录元宝...")
|
||
|
||
# 等待登录完成
|
||
await self._wait_for_login()
|
||
|
||
self.controller = YuanbaoController(self.page)
|
||
|
||
# 发送系统提示词（作为第一条消息）
|
||
print("[Agent] 发送系统提示词...")
|
||
await self._send_system_prompt()
|
||
|
||
print("[Agent] 开始监听对话...")
|
||
await self._conversation_loop()
|
||
|
||
async def _wait_for_login(self):
|
||
"""等待用户登录"""
|
||
print("[Agent] 登录完成后按回车继续...")
|
||
loop = asyncio.get_event_loop()
|
||
await loop.run_in_executor(None, input)
|
||
|
||
async def _send_system_prompt(self):
|
||
"""发送系统提示词"""
|
||
# 使用特殊前缀让元宝知道这是系统指令
|
||
prompt = f"""[系统初始化]
|
||
|
||
{SYSTEM_PROMPT}
|
||
|
||
---
|
||
现在你可以开始帮助用户了。当需要操作本地文件或执行命令时，请使用工具调用标记。"""
|
||
|
||
await self.controller.send_message(prompt)
|
||
print("[Agent] 系统提示词已发送")
|
||
|
||
# 等待元宝确认
|
||
await asyncio.sleep(3)
|
||
|
||
async def _conversation_loop(self):
|
||
"""对话主循环"""
|
||
last_ai_msg = None
|
||
|
||
while True:
|
||
try:
|
||
# 等待新的 AI 消息
|
||
print("[Agent] 等待元宝回复...")
|
||
ai_msg = await self.controller.wait_for_new_ai_message(last_ai_msg, timeout=180)
|
||
|
||
if not ai_msg:
|
||
print("[Agent] 等待超时，继续监听...")
|
||
continue
|
||
|
||
last_ai_msg = ai_msg
|
||
print(f"[Agent] 收到元宝消息 ({len(ai_msg)} 字符)")
|
||
|
||
# 检查是否包含工具调用
|
||
tool_call = self.controller.parse_tool_call(ai_msg)
|
||
|
||
if tool_call:
|
||
print(f"[Agent] 检测到工具调用: {tool_call.tool}")
|
||
|
||
# 执行工具
|
||
result = registry.execute(tool_call.tool, tool_call.args)
|
||
|
||
# 构造结果消息
|
||
result_msg = f"""[工具执行结果]
|
||
工具: {tool_call.tool}
|
||
参数: {tool_call.args}
|
||
成功: {result.success}
|
||
|
||
{result.content}
|
||
|
||
请基于以上结果继续。"""
|
||
|
||
print(f"[Agent] 工具执行完成，返回结果...")
|
||
|
||
# 发送结果回元宝
|
||
await asyncio.sleep(1)
|
||
await self.controller.send_message(result_msg)
|
||
|
||
# 更新 last_ai_msg 避免重复处理
|
||
last_ai_msg = None
|
||
else:
|
||
# 普通回复，记录到历史
|
||
self.conversation_history.append({
|
||
"role": "assistant",
|
||
"content": ai_msg
|
||
})
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n[Agent] 用户中断")
|
||
break
|
||
except Exception as e:
|
||
print(f"[Agent] 错误: {e}")
|
||
await asyncio.sleep(3)
|
||
|
||
|
||
async def main():
|
||
agent = YuanbaoAgent()
|
||
try:
|
||
await agent.run()
|
||
except Exception as e:
|
||
print(f"[Agent] 启动失败: {e}")
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|