Agent/examples/tools_examples.py

"""
工具系统完整示例

本文件展示 @tool 装饰器的所有用法，包括：

## 基础功能
1. 最简形式
2. 带 i18n 展示信息
3. 带可编辑参数
4. 需要用户确认
5. 带 context 参数
6. 同步工具
7. 复杂返回类型

## 高级功能
8. 域名过滤（URL Patterns）
9. 敏感数据处理（<secret> 占位符 + TOTP）
10. 工具使用统计
11. 组合所有功能

注意：
- uid 参数会由框架自动注入，不需要用户传递
- context 参数用于传递额外上下文（如浏览器会话、当前 URL 等）
- 返回值可以是字符串、字典或 ToolResult
"""

import asyncio
import json
from typing import List, Dict, Any, Optional
from agent import tool, ToolResult, ToolContext, get_tool_registry


# ============================================================
# 基础功能示例
# ============================================================

# 1. 最简形式
@tool()
async def hello_world(name: str, uid: str = "") -> Dict[str, str]:
	"""
	最简单的工具示例

	Args:
		name: 要问候的名字
		uid: 用户ID（自动注入）

	Returns:
		包含问候语的字典
	"""
	return {"greeting": f"Hello, {name}!"}


# 2. 带 i18n 展示信息的工具
@tool(
	display={
		"zh": {
			"name": "搜索内容",
			"params": {
				"query": "搜索关键词",
				"limit": "返回数量"
			}
		},
		"en": {
			"name": "Search Content",
			"params": {
				"query": "Search query",
				"limit": "Number of results"
			}
		}
	}
)
async def search_content(
	query: str,
	limit: int = 10,
	uid: str = ""
) -> List[Dict[str, Any]]:
	"""
	搜索用户的内容

	使用语义搜索查找相关内容。display 参数用于前端展示：
	- 工具名称会根据用户语言显示为"搜索内容"或"Search Content"
	- 参数名称也会相应翻译

	Args:
		query: 搜索查询文本
		limit: 返回结果数量（默认10）
		uid: 用户ID（自动注入）

	Returns:
		搜索结果列表，每个包含 id, title, content, score
	"""
	# 实际实现中会调用向量搜索
	return [
		{
			"id": "doc_001",
			"title": f"关于 {query} 的文档",
			"content": f"这是与 {query} 相关的内容...",
			"score": 0.95
		}
	]


# 3. 带可编辑参数的工具
@tool(
	editable_params=["query", "filters"],
	display={
		"zh": {
			"name": "高级搜索",
			"params": {
				"query": "搜索关键词",
				"filters": "过滤条件",
				"sort_by": "排序方式"
			}
		}
	}
)
async def advanced_search(
	query: str,
	filters: Optional[Dict[str, Any]] = None,
	sort_by: str = "relevance",
	limit: int = 20,
	uid: str = ""
) -> Dict[str, Any]:
	"""
	高级搜索工具（允许用户编辑参数）

	editable_params 指定哪些参数允许用户在 LLM 生成后编辑：
	- LLM 会先生成 query 和 filters
	- 用户可以在确认前修改这些参数
	- 适用于搜索、创建等需要用户微调的场景

	Args:
		query: 搜索查询
		filters: 过滤条件（如 {"type": "note", "date_range": "7d"}）
		sort_by: 排序方式（relevance/date/title）
		limit: 返回数量
		uid: 用户ID（自动注入）

	Returns:
		搜索结果和元数据
	"""
	return {
		"results": [
			{"id": "1", "title": "Result 1", "score": 0.9},
			{"id": "2", "title": "Result 2", "score": 0.8},
		],
		"total": 42,
		"query": query,
		"filters_applied": filters or {},
		"sort_by": sort_by
	}


# 4. 需要用户确认的危险操作
@tool(
	requires_confirmation=True,
	display={
		"zh": {
			"name": "删除内容",
			"params": {
				"content_id": "内容ID",
				"permanent": "永久删除"
			}
		}
	}
)
async def delete_content(
	content_id: str,
	permanent: bool = False,
	uid: str = ""
) -> Dict[str, Any]:
	"""
	删除内容（需要用户确认）

	requires_confirmation=True 表示这是一个危险操作：
	- LLM 调用此工具时，不会立即执行
	- 会先向用户展示操作详情，等待确认
	- 用户确认后才会真正执行

	适用场景：删除操作、发送消息、修改重要设置、任何不可逆操作

	Args:
		content_id: 要删除的内容ID
		permanent: 是否永久删除（False=移到回收站）
		uid: 用户ID（自动注入）

	Returns:
		删除结果
	"""
	return {
		"success": True,
		"content_id": content_id,
		"permanent": permanent,
		"message": f"内容 {content_id} 已{'永久删除' if permanent else '移到回收站'}"
	}


# 5. 带 context 参数的工具
@tool(
	display={
		"zh": {"name": "获取相关推荐", "params": {"top_k": "推荐数量"}}
	}
)
async def get_recommendations(
	top_k: int = 5,
	uid: str = "",
	context: Optional[Dict[str, Any]] = None
) -> List[Dict[str, Any]]:
	"""
	获取相关推荐（使用 context 获取额外信息）

	context 参数用于传递执行上下文，由框架自动注入：
	- 当前阅读位置 (current_location)
	- 当前会话 ID (session_id)
	- 排除的内容 ID (exclude_ids)

	Args:
		top_k: 返回推荐数量
		uid: 用户ID（自动注入）
		context: 执行上下文（自动注入）

	Returns:
		推荐列表
	"""
	current_location = None
	if context:
		current_location = context.get("current_location")

	return [
		{
			"id": "rec_001",
			"title": "推荐内容 1",
			"reason": f"基于当前位置 {current_location}" if current_location else "基于您的兴趣"
		}
	]


# 6. 同步工具（非 async）
@tool()
def format_text(
	text: str,
	format_type: str = "markdown",
	uid: str = ""
) -> str:
	"""
	格式化文本（同步工具）

	不需要 async 的工具可以定义为普通函数。
	框架会自动检测并正确调用。

	适用于：纯计算操作、文本处理、不需要 I/O 的操作

	Args:
		text: 要格式化的文本
		format_type: 格式类型（markdown/plain/html）
		uid: 用户ID（自动注入）

	Returns:
		格式化后的文本
	"""
	if format_type == "markdown":
		return f"**{text}**"
	elif format_type == "html":
		return f"<p>{text}</p>"
	else:
		return text


# 7. 使用 ToolResult 的工具
@tool()
async def analyze_content(
	content_id: str,
	analysis_types: Optional[List[str]] = None,
	uid: str = ""
) -> ToolResult:
	"""
	分析内容（使用 ToolResult）

	ToolResult 支持双层记忆管理：
	- output: 完整结果（可能很长）
	- long_term_memory: 简短摘要（永久保存）

	Args:
		content_id: 要分析的内容ID
		analysis_types: 分析类型列表（sentiment/keywords/summary）
		uid: 用户ID（自动注入）

	Returns:
		ToolResult 包含分析结果
	"""
	types = analysis_types or ["sentiment", "keywords"]

	result = {
		"content_id": content_id,
		"analyses": {}
	}

	if "sentiment" in types:
		result["analyses"]["sentiment"] = {
			"score": 0.8,
			"label": "positive",
			"confidence": 0.92
		}

	if "keywords" in types:
		result["analyses"]["keywords"] = [
			{"word": "AI", "weight": 0.9},
			{"word": "学习", "weight": 0.7}
		]

	return ToolResult(
		title=f"Analysis of {content_id}",
		output=json.dumps(result, indent=2, ensure_ascii=False),
		long_term_memory=f"Analyzed {content_id}: {', '.join(types)}",
		metadata={"types": types}
	)


# ============================================================
# 高级功能示例
# ============================================================

# 8. 域名过滤示例
@tool(url_patterns=["*.google.com", "www.google.*"])
async def google_search(query: str, uid: str = "") -> ToolResult:
	"""
	Google 搜索（仅在 Google 页面可用）

	使用 url_patterns 限制工具只在特定域名显示。
	在 Google 页面时，此工具会出现在可用工具列表中。
	在其他页面时，此工具会被过滤掉。

	Args:
		query: 搜索查询
		uid: 用户ID（自动注入）

	Returns:
		搜索结果
	"""
	return ToolResult(
		title="Google Search",
		output=f"Searching Google for: {query}",
		long_term_memory=f"Searched Google for '{query}'"
	)


@tool(url_patterns=["*.github.com"])
async def create_github_issue(
	title: str,
	body: str,
	uid: str = ""
) -> ToolResult:
	"""
	创建 GitHub Issue（仅在 GitHub 页面可用）

	Args:
		title: Issue 标题
		body: Issue 内容
		uid: 用户ID（自动注入）

	Returns:
		创建结果
	"""
	return ToolResult(
		title="Issue Created",
		output=f"Created issue: {title}",
		long_term_memory=f"Created GitHub issue: {title}"
	)


@tool()  # 无 url_patterns，所有页面都可用
async def take_screenshot(uid: str = "") -> ToolResult:
	"""截图（所有页面都可用）"""
	return ToolResult(
		title="Screenshot",
		output="Screenshot taken",
		attachments=["screenshot_001.png"]
	)


# 9. 敏感数据处理示例
@tool(url_patterns=["*.github.com"])
async def github_login(
	username: str,
	password: str,
	totp_code: str,
	uid: str = ""
) -> ToolResult:
	"""
	GitHub 登录（支持敏感数据占位符）

	LLM 会输出类似：
	{
		"username": "user@example.com",
		"password": "<secret>github_password</secret>",
		"totp_code": "<secret>github_2fa_bu_2fa_code</secret>"
	}

	执行时会自动替换为实际值。

	Args:
		username: 用户名
		password: 密码（可以是占位符）
		totp_code: TOTP 验证码（可以是占位符，自动生成）
		uid: 用户ID（自动注入）

	Returns:
		登录结果
	"""
	# 注意：password 和 totp_code 在到达这里时已经被替换
	return ToolResult(
		title="Login Successful",
		output=f"Logged in as {username}",
		long_term_memory=f"Logged in to GitHub as {username}"
	)


# 10. 组合所有功能
@tool(
	url_patterns=["*.example.com"],
	requires_confirmation=True,
	editable_params=["message"],
	display={
		"zh": {
			"name": "发送认证消息",
			"params": {
				"recipient": "接收者",
				"message": "消息内容",
				"api_key": "API密钥"
			}
		}
	}
)
async def send_authenticated_message(
	recipient: str,
	message: str,
	api_key: str,
	ctx: ToolContext,
	uid: str = ""
) -> ToolResult:
	"""
	发送消息（组合多个功能）

	展示所有高级功能：
	- 仅在 example.com 可用（域名过滤）
	- 需要用户确认（危险操作）
	- 消息可编辑（用户微调）
	- API key 使用敏感数据占位符
	- 使用 ToolContext 获取上下文

	Args:
		recipient: 接收者
		message: 消息内容
		api_key: API密钥（可以是占位符）
		ctx: 工具上下文
		uid: 用户ID（自动注入）

	Returns:
		发送结果
	"""
	# api_key 会从 <secret>api_key</secret> 替换为实际值
	# ctx 包含 page_url, browser_session 等信息

	return ToolResult(
		title="Message Sent",
		output=f"Sent to {recipient}: {message}",
		long_term_memory=f"Sent message to {recipient} on {ctx.page_url}",
		metadata={"recipient": recipient}
	)


# ============================================================
# 使用示例
# ============================================================

async def main():
	registry = get_tool_registry()

	print("=" * 60)
	print("工具系统完整示例")
	print("=" * 60)

	# ============================================================
	# 示例 1：基础工具调用
	# ============================================================
	print("\n1. 基础工具调用")
	print("-" * 60)

	result = await registry.execute("hello_world", {"name": "Alice"})
	print(f"hello_world: {result}")

	result = await registry.execute("search_content", {"query": "Python", "limit": 5})
	print(f"search_content: {result}")

	# ============================================================
	# 示例 2：域名过滤
	# ============================================================
	print("\n\n2. 域名过滤示例")
	print("-" * 60)

	# 在 Google 页面
	google_url = "https://www.google.com/search?q=test"
	google_tools = registry.get_tool_names(google_url)
	print(f"在 {google_url} 可用的工具:")
	print(f"  包含 google_search: {'google_search' in google_tools}")

	# 在 GitHub 页面
	github_url = "https://github.com/user/repo"
	github_tools = registry.get_tool_names(github_url)
	print(f"\n在 {github_url} 可用的工具:")
	print(f"  包含 create_github_issue: {'create_github_issue' in github_tools}")
	print(f"  包含 google_search: {'google_search' in github_tools}")

	# ============================================================
	# 示例 3：敏感数据处理
	# ============================================================
	print("\n\n3. 敏感数据处理示例")
	print("-" * 60)

	# 配置敏感数据
	sensitive_data = {
		"*.github.com": {
			"github_password": "my_secret_password",
			"github_2fa_bu_2fa_code": "JBSWY3DPEHPK3PXP"  # TOTP secret
		}
	}

	# 模拟 LLM 输出（包含占位符）
	llm_output_args = {
		"username": "user@example.com",
		"password": "<secret>github_password</secret>",
		"totp_code": "<secret>github_2fa_bu_2fa_code</secret>"
	}

	print("LLM 输出的参数（包含占位符）:")
	print(f"  {llm_output_args}")

	# 执行工具（自动替换敏感数据）
	result = await registry.execute(
		"github_login",
		llm_output_args,
		context={"page_url": "https://github.com/login"},
		sensitive_data=sensitive_data
	)

	print(f"\n执行结果（密码已替换）:")
	print(f"  {result}")

	# ============================================================
	# 示例 4：工具统计
	# ============================================================
	print("\n\n4. 工具统计示例")
	print("-" * 60)

	# 模拟多次调用
	for i in range(5):
		await registry.execute("google_search", {"query": f"test {i}"})

	await registry.execute("take_screenshot", {})
	await registry.execute("take_screenshot", {})

	# 查看统计
	stats = registry.get_stats()
	print("工具使用统计:")
	for tool_name, tool_stats in stats.items():
		if tool_stats["call_count"] > 0:
			print(f"\n  {tool_name}:")
			print(f"    调用次数: {tool_stats['call_count']}")
			print(f"    成功率: {tool_stats['success_rate']:.1%}")
			print(f"    平均执行时间: {tool_stats['average_duration']:.3f}s")

	# 获取 Top 工具
	print("\n\nTop 3 最常用工具:")
	top_tools = registry.get_top_tools(limit=3, by="call_count")
	for i, tool_name in enumerate(top_tools, 1):
		tool_stats = stats[tool_name]
		print(f"  {i}. {tool_name} ({tool_stats['call_count']} 次调用)")


if __name__ == "__main__":
	asyncio.run(main())