Spaces:

pzp5700
/

Paper2Any

Sleeping

Paper2Any / dataflow_agent /trajectory /models.py

HF Space (clean): Paper2PPT + PPT Polish

d82bbe4 10 months ago

12.1 kB

"""
TRJ 数据模型定义
定义标准的轨迹数据结构，支持：
1. ReAct 模式：Context -> Thought -> Action -> Observation
2. Workflow 模式：State_In -> Node Processing -> State_Update
"""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional
import uuid


class TrajectoryMode(str, Enum):
    """轨迹模式"""
    REACT = "react"
    WORKFLOW = "workflow"
    HYBRID = "hybrid"  # 混合模式


class StepRole(str, Enum):
    """步骤角色"""
    AGENT = "agent"
    ENVIRONMENT = "environment"
    SYSTEM_NODE = "system_node"
    TOOL = "tool"
    USER = "user"


class ActionType(str, Enum):
    """动作类型"""
    TOOL_CALL = "tool_call"
    RESPONSE = "response"
    STATE_UPDATE = "state_update"
    LLM_CALL = "llm_call"
    MULTIMODAL = "multimodal"


@dataclass
class ToolCallRecord:
    """工具调用记录"""
    tool_name: str
    tool_args: Dict[str, Any]
    tool_result: Any
    timestamp: str
    duration_ms: Optional[float] = None
    error: Optional[str] = None


@dataclass
class LLMCallRecord:
    """LLM 调用记录"""
    model: str
    messages_in: List[Dict[str, Any]]  # 输入消息
    response: str  # 输出响应
    timestamp: str
    duration_ms: Optional[float] = None
    token_usage: Optional[Dict[str, int]] = None  # {"prompt": x, "completion": y}
    temperature: Optional[float] = None


@dataclass
class MultimodalData:
    """多模态数据"""
    type: str  # "image" | "audio" | "video"
    path: Optional[str] = None  # 文件路径
    url: Optional[str] = None  # URL
    base64: Optional[str] = None  # Base64 编码（不推荐存储大数据）
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class TrajectoryStep:
    """
    单个执行步骤
    
    对于 ReAct 模式：
    - input_context: Agent 看到的上下文
    - thought: Agent 的思考过程
    - action_type: 动作类型（tool_call/response）
    - action_payload: 动作内容
    - observation: 环境反馈
    
    对于 Workflow 模式：
    - input_context: 节点输入状态
    - node_output: 节点输出/状态更新
    """
    step_index: int
    node_name: str
    role: str  # StepRole 的值
    timestamp: str
    
    # 输入上下文
    input_context: Dict[str, Any] = field(default_factory=dict)
    
    # ReAct 特有字段
    thought: Optional[str] = None
    action_type: Optional[str] = None  # ActionType 的值
    action_payload: Optional[Dict[str, Any]] = None
    observation: Optional[str] = None
    
    # 通用输出
    node_output: Optional[Dict[str, Any]] = None
    
    # 详细记录
    llm_calls: List[LLMCallRecord] = field(default_factory=list)
    tool_calls: List[ToolCallRecord] = field(default_factory=list)
    
    # 多模态数据
    multimodal_input: Optional[MultimodalData] = None
    multimodal_output: Optional[MultimodalData] = None
    
    # 错误信息
    error: Optional[str] = None
    
    # 执行时间
    duration_ms: Optional[float] = None
    
    # 额外元数据
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        result = {
            "step_index": self.step_index,
            "node_name": self.node_name,
            "role": self.role,
            "timestamp": self.timestamp,
        }
        
        # 只添加非空字段
        if self.input_context:
            result["input_context"] = self.input_context
        if self.thought:
            result["thought"] = self.thought
        if self.action_type:
            result["action_type"] = self.action_type
        if self.action_payload:
            result["action_payload"] = self.action_payload
        if self.observation:
            result["observation"] = self.observation
        if self.node_output:
            result["node_output"] = self.node_output
        if self.llm_calls:
            result["llm_calls"] = [self._llm_call_to_dict(c) for c in self.llm_calls]
        if self.tool_calls:
            result["tool_calls"] = [self._tool_call_to_dict(c) for c in self.tool_calls]
        if self.multimodal_input:
            result["multimodal_input"] = self._multimodal_to_dict(self.multimodal_input)
        if self.multimodal_output:
            result["multimodal_output"] = self._multimodal_to_dict(self.multimodal_output)
        if self.error:
            result["error"] = self.error
        if self.duration_ms is not None:
            result["duration_ms"] = self.duration_ms
        if self.metadata:
            result["metadata"] = self.metadata
            
        return result
    
    @staticmethod
    def _llm_call_to_dict(call: LLMCallRecord) -> Dict[str, Any]:
        return {
            "model": call.model,
            "messages_in": call.messages_in,
            "response": call.response,
            "timestamp": call.timestamp,
            "duration_ms": call.duration_ms,
            "token_usage": call.token_usage,
            "temperature": call.temperature,
        }
    
    @staticmethod
    def _tool_call_to_dict(call: ToolCallRecord) -> Dict[str, Any]:
        return {
            "tool_name": call.tool_name,
            "tool_args": call.tool_args,
            "tool_result": call.tool_result,
            "timestamp": call.timestamp,
            "duration_ms": call.duration_ms,
            "error": call.error,
        }
    
    @staticmethod
    def _multimodal_to_dict(data: MultimodalData) -> Dict[str, Any]:
        result = {"type": data.type}
        if data.path:
            result["path"] = data.path
        if data.url:
            result["url"] = data.url
        if data.metadata:
            result["metadata"] = data.metadata
        # 不导出 base64 以减小文件大小
        return result


@dataclass
class TrajectoryFeedback:
    """用户反馈"""
    score: Optional[int] = None  # 1-5 评分
    comment: Optional[str] = None
    edited_response: Optional[str] = None  # 用户修改后的回答（用于 SFT）
    labels: List[str] = field(default_factory=list)  # 标签，如 ["good", "accurate"]
    timestamp: Optional[str] = None


@dataclass
class Trajectory:
    """
    完整的执行轨迹
    
    包含三个核心部分：
    1. Metadata: 元数据
    2. Steps: 执行步骤列表
    3. Outcome: 最终结果和反馈
    """
    # ===== 元数据 =====
    trace_id: str
    workflow_name: str
    timestamp: str
    status: str  # "success" | "failed" | "partial"
    mode: str  # TrajectoryMode 的值
    
    # 可选元数据
    user_id: Optional[str] = None
    session_id: Optional[str] = None
    version: str = "1.0"
    
    # ===== 输入 =====
    inputs: Dict[str, Any] = field(default_factory=dict)
    
    # ===== 执行步骤 =====
    steps: List[TrajectoryStep] = field(default_factory=list)
    
    # ===== 输出 =====
    final_output: Any = None
    
    # ===== 反馈 =====
    feedback: Optional[TrajectoryFeedback] = None
    
    # ===== 统计信息 =====
    total_duration_ms: Optional[float] = None
    total_llm_calls: int = 0
    total_tool_calls: int = 0
    total_tokens: Optional[Dict[str, int]] = None
    
    # ===== 额外元数据 =====
    metadata: Dict[str, Any] = field(default_factory=dict)

    @staticmethod
    def generate_trace_id() -> str:
        """生成唯一的 trace_id"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = uuid.uuid4().hex[:8]
        return f"trj_{timestamp}_{unique_id}"

    def add_step(self, step: TrajectoryStep):
        """添加执行步骤"""
        self.steps.append(step)
        # 更新统计
        self.total_llm_calls += len(step.llm_calls)
        self.total_tool_calls += len(step.tool_calls)

    def set_feedback(self, score: int = None, comment: str = None, 
                     edited_response: str = None, labels: List[str] = None):
        """设置用户反馈"""
        self.feedback = TrajectoryFeedback(
            score=score,
            comment=comment,
            edited_response=edited_response,
            labels=labels or [],
            timestamp=datetime.now().isoformat()
        )

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典（用于 JSON 导出）"""
        result = {
            # 元数据
            "trace_id": self.trace_id,
            "workflow_name": self.workflow_name,
            "timestamp": self.timestamp,
            "status": self.status,
            "mode": self.mode,
            "version": self.version,
            
            # 输入
            "inputs": self.inputs,
            
            # 步骤
            "steps": [step.to_dict() for step in self.steps],
            
            # 输出
            "final_output": self.final_output,
            
            # 统计
            "statistics": {
                "total_steps": len(self.steps),
                "total_llm_calls": self.total_llm_calls,
                "total_tool_calls": self.total_tool_calls,
                "total_duration_ms": self.total_duration_ms,
                "total_tokens": self.total_tokens,
            }
        }
        
        # 可选字段
        if self.user_id:
            result["user_id"] = self.user_id
        if self.session_id:
            result["session_id"] = self.session_id
        if self.feedback:
            result["feedback"] = {
                "score": self.feedback.score,
                "comment": self.feedback.comment,
                "edited_response": self.feedback.edited_response,
                "labels": self.feedback.labels,
                "timestamp": self.feedback.timestamp,
            }
        if self.metadata:
            result["metadata"] = self.metadata
            
        return result

    def to_sft_format(self) -> List[Dict[str, str]]:
        """
        转换为 SFT 训练格式（OpenAI messages 格式）
        
        Returns:
            [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
        """
        messages = []
        
        for step in self.steps:
            if step.role == StepRole.AGENT.value:
                # Agent 的输出
                content = ""
                if step.thought:
                    content += f"<thought>{step.thought}</thought>\n"
                if step.action_type == ActionType.TOOL_CALL.value and step.action_payload:
                    tool_name = step.action_payload.get("tool_name", "")
                    tool_args = step.action_payload.get("tool_args", "")
                    content += f"<call>{tool_name}({tool_args})</call>"
                elif step.node_output:
                    content += str(step.node_output)
                
                if content:
                    messages.append({"role": "assistant", "content": content})
                    
            elif step.role in [StepRole.ENVIRONMENT.value, StepRole.TOOL.value]:
                # 工具/环境的输出
                if step.observation:
                    messages.append({"role": "tool", "content": step.observation})
                    
            elif step.role == StepRole.USER.value:
                # 用户输入
                if step.input_context:
                    content = step.input_context.get("query", str(step.input_context))
                    messages.append({"role": "user", "content": content})
        
        return messages

    def to_dpo_format(self) -> Dict[str, Any]:
        """
        转换为 DPO 训练格式
        
        Returns:
            {"prompt": "...", "chosen": [...], "rejected": [...]}
        """
        # 提取 prompt
        prompt = self.inputs.get("query", self.inputs.get("target", ""))
        
        # 当前轨迹作为 chosen 或 rejected
        trajectory_steps = self.to_sft_format()
        
        return {
            "prompt": prompt,
            "trajectory": trajectory_steps,
            "score": self.feedback.score if self.feedback else None,
            "status": self.status,
        }

	"""
	TRJ 数据模型定义

	定义标准的轨迹数据结构，支持：
	1. ReAct 模式：Context -> Thought -> Action -> Observation
	2. Workflow 模式：State_In -> Node Processing -> State_Update
	"""

	from dataclasses import dataclass, field
	from datetime import datetime
	from enum import Enum
	from typing import Any, Dict, List, Optional
	import uuid


	class TrajectoryMode(str, Enum):
	"""轨迹模式"""
	REACT = "react"
	WORKFLOW = "workflow"
	HYBRID = "hybrid" # 混合模式


	class StepRole(str, Enum):
	"""步骤角色"""
	AGENT = "agent"
	ENVIRONMENT = "environment"
	SYSTEM_NODE = "system_node"
	TOOL = "tool"
	USER = "user"


	class ActionType(str, Enum):
	"""动作类型"""
	TOOL_CALL = "tool_call"
	RESPONSE = "response"
	STATE_UPDATE = "state_update"
	LLM_CALL = "llm_call"
	MULTIMODAL = "multimodal"


	@dataclass
	class ToolCallRecord:
	"""工具调用记录"""
	tool_name: str
	tool_args: Dict[str, Any]
	tool_result: Any
	timestamp: str
	duration_ms: Optional[float] = None
	error: Optional[str] = None


	@dataclass
	class LLMCallRecord:
	"""LLM 调用记录"""
	model: str
	messages_in: List[Dict[str, Any]] # 输入消息
	response: str # 输出响应
	timestamp: str
	duration_ms: Optional[float] = None
	token_usage: Optional[Dict[str, int]] = None # {"prompt": x, "completion": y}
	temperature: Optional[float] = None


	@dataclass
	class MultimodalData:
	"""多模态数据"""
	type: str # "image" \| "audio" \| "video"
	path: Optional[str] = None # 文件路径
	url: Optional[str] = None # URL
	base64: Optional[str] = None # Base64 编码（不推荐存储大数据）
	metadata: Dict[str, Any] = field(default_factory=dict)


	@dataclass
	class TrajectoryStep:
	"""
	单个执行步骤

	对于 ReAct 模式：
	- input_context: Agent 看到的上下文
	- thought: Agent 的思考过程
	- action_type: 动作类型（tool_call/response）
	- action_payload: 动作内容
	- observation: 环境反馈

	对于 Workflow 模式：
	- input_context: 节点输入状态
	- node_output: 节点输出/状态更新
	"""
	step_index: int
	node_name: str
	role: str # StepRole 的值
	timestamp: str

	# 输入上下文
	input_context: Dict[str, Any] = field(default_factory=dict)

	# ReAct 特有字段
	thought: Optional[str] = None
	action_type: Optional[str] = None # ActionType 的值
	action_payload: Optional[Dict[str, Any]] = None
	observation: Optional[str] = None

	# 通用输出
	node_output: Optional[Dict[str, Any]] = None

	# 详细记录
	llm_calls: List[LLMCallRecord] = field(default_factory=list)
	tool_calls: List[ToolCallRecord] = field(default_factory=list)

	# 多模态数据
	multimodal_input: Optional[MultimodalData] = None
	multimodal_output: Optional[MultimodalData] = None

	# 错误信息
	error: Optional[str] = None

	# 执行时间
	duration_ms: Optional[float] = None

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

	def to_dict(self) -> Dict[str, Any]:
	"""转换为字典"""
	result = {
	"step_index": self.step_index,
	"node_name": self.node_name,
	"role": self.role,
	"timestamp": self.timestamp,
	}

	# 只添加非空字段
	if self.input_context:
	result["input_context"] = self.input_context
	if self.thought:
	result["thought"] = self.thought
	if self.action_type:
	result["action_type"] = self.action_type
	if self.action_payload:
	result["action_payload"] = self.action_payload
	if self.observation:
	result["observation"] = self.observation
	if self.node_output:
	result["node_output"] = self.node_output
	if self.llm_calls:
	result["llm_calls"] = [self._llm_call_to_dict(c) for c in self.llm_calls]
	if self.tool_calls:
	result["tool_calls"] = [self._tool_call_to_dict(c) for c in self.tool_calls]
	if self.multimodal_input:
	result["multimodal_input"] = self._multimodal_to_dict(self.multimodal_input)
	if self.multimodal_output:
	result["multimodal_output"] = self._multimodal_to_dict(self.multimodal_output)
	if self.error:
	result["error"] = self.error
	if self.duration_ms is not None:
	result["duration_ms"] = self.duration_ms
	if self.metadata:
	result["metadata"] = self.metadata

	return result

	@staticmethod
	def _llm_call_to_dict(call: LLMCallRecord) -> Dict[str, Any]:
	return {
	"model": call.model,
	"messages_in": call.messages_in,
	"response": call.response,
	"timestamp": call.timestamp,
	"duration_ms": call.duration_ms,
	"token_usage": call.token_usage,
	"temperature": call.temperature,
	}

	@staticmethod
	def _tool_call_to_dict(call: ToolCallRecord) -> Dict[str, Any]:
	return {
	"tool_name": call.tool_name,
	"tool_args": call.tool_args,
	"tool_result": call.tool_result,
	"timestamp": call.timestamp,
	"duration_ms": call.duration_ms,
	"error": call.error,
	}

	@staticmethod
	def _multimodal_to_dict(data: MultimodalData) -> Dict[str, Any]:
	result = {"type": data.type}
	if data.path:
	result["path"] = data.path
	if data.url:
	result["url"] = data.url
	if data.metadata:
	result["metadata"] = data.metadata
	# 不导出 base64 以减小文件大小
	return result


	@dataclass
	class TrajectoryFeedback:
	"""用户反馈"""
	score: Optional[int] = None # 1-5 评分
	comment: Optional[str] = None
	edited_response: Optional[str] = None # 用户修改后的回答（用于 SFT）
	labels: List[str] = field(default_factory=list) # 标签，如 ["good", "accurate"]
	timestamp: Optional[str] = None


	@dataclass
	class Trajectory:
	"""
	完整的执行轨迹

	包含三个核心部分：
	1. Metadata: 元数据
	2. Steps: 执行步骤列表
	3. Outcome: 最终结果和反馈
	"""
	# ===== 元数据 =====
	trace_id: str
	workflow_name: str
	timestamp: str
	status: str # "success" \| "failed" \| "partial"
	mode: str # TrajectoryMode 的值

	# 可选元数据
	user_id: Optional[str] = None
	session_id: Optional[str] = None
	version: str = "1.0"

	# ===== 输入 =====
	inputs: Dict[str, Any] = field(default_factory=dict)

	# ===== 执行步骤 =====
	steps: List[TrajectoryStep] = field(default_factory=list)

	# ===== 输出 =====
	final_output: Any = None

	# ===== 反馈 =====
	feedback: Optional[TrajectoryFeedback] = None

	# ===== 统计信息 =====
	total_duration_ms: Optional[float] = None
	total_llm_calls: int = 0
	total_tool_calls: int = 0
	total_tokens: Optional[Dict[str, int]] = None

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

	@staticmethod
	def generate_trace_id() -> str:
	"""生成唯一的 trace_id"""
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	unique_id = uuid.uuid4().hex[:8]
	return f"trj_{timestamp}_{unique_id}"

	def add_step(self, step: TrajectoryStep):
	"""添加执行步骤"""
	self.steps.append(step)
	# 更新统计
	self.total_llm_calls += len(step.llm_calls)
	self.total_tool_calls += len(step.tool_calls)

	def set_feedback(self, score: int = None, comment: str = None,
	edited_response: str = None, labels: List[str] = None):
	"""设置用户反馈"""
	self.feedback = TrajectoryFeedback(
	score=score,
	comment=comment,
	edited_response=edited_response,
	labels=labels or [],
	timestamp=datetime.now().isoformat()
	)

	def to_dict(self) -> Dict[str, Any]:
	"""转换为字典（用于 JSON 导出）"""
	result = {
	# 元数据
	"trace_id": self.trace_id,
	"workflow_name": self.workflow_name,
	"timestamp": self.timestamp,
	"status": self.status,
	"mode": self.mode,
	"version": self.version,

	# 输入
	"inputs": self.inputs,

	# 步骤
	"steps": [step.to_dict() for step in self.steps],

	# 输出
	"final_output": self.final_output,

	# 统计
	"statistics": {
	"total_steps": len(self.steps),
	"total_llm_calls": self.total_llm_calls,
	"total_tool_calls": self.total_tool_calls,
	"total_duration_ms": self.total_duration_ms,
	"total_tokens": self.total_tokens,
	}
	}

	# 可选字段
	if self.user_id:
	result["user_id"] = self.user_id
	if self.session_id:
	result["session_id"] = self.session_id
	if self.feedback:
	result["feedback"] = {
	"score": self.feedback.score,
	"comment": self.feedback.comment,
	"edited_response": self.feedback.edited_response,
	"labels": self.feedback.labels,
	"timestamp": self.feedback.timestamp,
	}
	if self.metadata:
	result["metadata"] = self.metadata

	return result

	def to_sft_format(self) -> List[Dict[str, str]]:
	"""
	转换为 SFT 训练格式（OpenAI messages 格式）

	Returns:
	[{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
	"""
	messages = []

	for step in self.steps:
	if step.role == StepRole.AGENT.value:
	# Agent 的输出
	content = ""
	if step.thought:
	content += f"<thought>{step.thought}</thought>\n"
	if step.action_type == ActionType.TOOL_CALL.value and step.action_payload:
	tool_name = step.action_payload.get("tool_name", "")
	tool_args = step.action_payload.get("tool_args", "")
	content += f"<call>{tool_name}({tool_args})</call>"
	elif step.node_output:
	content += str(step.node_output)

	if content:
	messages.append({"role": "assistant", "content": content})

	elif step.role in [StepRole.ENVIRONMENT.value, StepRole.TOOL.value]:
	# 工具/环境的输出
	if step.observation:
	messages.append({"role": "tool", "content": step.observation})

	elif step.role == StepRole.USER.value:
	# 用户输入
	if step.input_context:
	content = step.input_context.get("query", str(step.input_context))
	messages.append({"role": "user", "content": content})

	return messages

	def to_dpo_format(self) -> Dict[str, Any]:
	"""
	转换为 DPO 训练格式

	Returns:
	{"prompt": "...", "chosen": [...], "rejected": [...]}
	"""
	# 提取 prompt
	prompt = self.inputs.get("query", self.inputs.get("target", ""))

	# 当前轨迹作为 chosen 或 rejected
	trajectory_steps = self.to_sft_format()

	return {
	"prompt": prompt,
	"trajectory": trajectory_steps,
	"score": self.feedback.score if self.feedback else None,
	"status": self.status,
	}