mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-16 23:41:10 +00:00
63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
"""Trajectory preprocessing operation for task memory generation.
|
|
|
|
This module provides operations to preprocess and classify trajectories
|
|
into success and failure categories based on score thresholds.
|
|
"""
|
|
|
|
from typing import List, Dict
|
|
|
|
from flowllm.core.context import C
|
|
from flowllm.core.op import BaseAsyncOp
|
|
from loguru import logger
|
|
|
|
from reme_ai.schema import Trajectory
|
|
|
|
|
|
@C.register_op()
|
|
class TrajectoryPreprocessOp(BaseAsyncOp):
|
|
"""Preprocess trajectories: validate and classify by success/failure.
|
|
|
|
This operation classifies trajectories into success and failure categories
|
|
based on score thresholds, preparing them for downstream memory extraction
|
|
operations.
|
|
"""
|
|
|
|
file_path: str = __file__
|
|
|
|
async def async_execute(self):
|
|
"""Preprocess trajectories: validate and classify"""
|
|
trajectories: list = self.context.get("trajectories", [])
|
|
trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories]
|
|
|
|
# Classify trajectories
|
|
classified = self._classify_trajectories(trajectories)
|
|
logger.info(
|
|
f"Classified trajectories - Success: {len(classified['success'])}, "
|
|
f"Failure: {len(classified['failure'])}, All: {len(classified['all'])}",
|
|
)
|
|
|
|
# Set context for downstream operators
|
|
self.context.success_trajectories = classified["success"]
|
|
self.context.failure_trajectories = classified["failure"]
|
|
self.context.all_trajectories = classified["all"]
|
|
|
|
def _classify_trajectories(self, trajectories: List[Trajectory]) -> Dict[str, List[Trajectory]]:
|
|
"""Classify trajectories based on score threshold"""
|
|
success_trajectories = []
|
|
failure_trajectories = []
|
|
|
|
success_threshold = self.op_params.get("success_threshold", 1.0)
|
|
|
|
for traj in trajectories:
|
|
is_success = traj.score >= success_threshold
|
|
|
|
if is_success:
|
|
success_trajectories.append(traj)
|
|
else:
|
|
failure_trajectories.append(traj)
|
|
|
|
return {
|
|
"success": success_trajectories,
|
|
"failure": failure_trajectories,
|
|
"all": trajectories,
|
|
}
|