Agent/cyber_agent/tools/models.py

"""
Tool Models - 工具系统核心数据模型

定义：
1. ToolResult: 工具执行结果（支持双层记忆管理）
2. ToolContext: 工具执行上下文（依赖注入）
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol


@dataclass
class ToolResult:
	"""
	工具执行结果

	支持双层记忆管理（参考 Browser-Use 的 ActionResult）：
	- output: 主要输出，可能很长，可以配置只给 LLM 看一次
	- long_term_memory: 简短摘要，永久保存在对话历史中

	这种设计避免了大量临时内容占用 context。
	"""

	# 主要输出（临时内容）
	title: str  # 简短标题，用于展示
	output: str  # 主要输出内容

	# 记忆管理
	long_term_memory: Optional[str] = None  # 永久记忆（简短摘要）
	include_output_only_once: bool = False  # output 是否只给 LLM 看一次

	# 元数据
	metadata: Dict[str, Any] = field(default_factory=dict)

	# 状态标志
	truncated: bool = False  # 输出是否被截断
	error: Optional[str] = None  # 错误信息（如果执行失败）

	# 附件支持（用于浏览器自动化等场景）
	attachments: List[str] = field(default_factory=list)  # 文件路径列表
	images: List[Dict[str, Any]] = field(default_factory=list)  # 图片列表
	artifact_refs: List[Dict[str, Any]] = field(default_factory=list)  # Validator可解析的标准只读产物引用

	# Token追踪（用于工具内部LLM调用）
	tool_usage: Optional[Dict[str, Any]] = None  # 格式：{"model": "...", "prompt_tokens": 100, "completion_tokens": 50, "cost": 0.0}

	def to_llm_message(self, first_time: bool = True) -> str:
		"""
		转换为给 LLM 的消息

		Args:
			first_time: 是否第一次展示（影响 include_output_only_once 的行为）

		Returns:
			给 LLM 的消息字符串
		"""
		# 如果有错误，优先返回错误
		if self.error:
			return f"Error: {self.error}"

		# 构建消息
		parts = []

		# 标题
		if self.title:
			parts.append(f"# {self.title}")

		# 主要输出
		if first_time or not self.include_output_only_once:
			if self.output:
				parts.append(self.output)
				if self.truncated:
					parts.append("(Output truncated)")

		# 长期记忆（永远包含）
		if self.long_term_memory:
			parts.append(f"\nSummary: {self.long_term_memory}")

		# 附件信息
		if self.attachments:
			parts.append(f"\nAttachments: {', '.join(self.attachments)}")

		return "\n\n".join(parts)


class ToolContext(Protocol):
	"""
	工具执行上下文（依赖注入）

	工具函数可以声明需要哪些上下文字段，框架自动注入。

	使用 Protocol 允许不同实现提供不同的上下文字段。
	"""

	# 基础字段（所有工具都可用）
	trace_id: str
	step_id: str
	uid: Optional[str]

	# 浏览器相关（Browser-Use 集成）
	browser_session: Optional[Any]  # BrowserSession 实例
	page_url: Optional[str]  # 当前页面 URL
	file_system: Optional[Any]  # FileSystem 实例
	sensitive_data: Optional[Dict[str, Any]]  # 敏感数据字典

	# 其他可扩展字段
	context: Optional[Dict[str, Any]]  # 额外上下文数据


@dataclass
class ToolContextImpl:
	"""ToolContext 的默认实现"""

	# 基础字段
	trace_id: str
	step_id: str
	uid: Optional[str] = None

	# 浏览器相关
	browser_session: Optional[Any] = None
	page_url: Optional[str] = None
	file_system: Optional[Any] = None
	sensitive_data: Optional[Dict[str, Any]] = None

	# 额外上下文
	context: Optional[Dict[str, Any]] = None