ReMe/reme_ai/summary/task/memory_validation_op.py
jinli.yl 0c2fb9a522 refactor(reme_ai): formatting and minor code improvements
- Removed unnecessary comments
- Improved code formatting across multiple files
- Updated import statements and removed unused imports
- Refactored some function definitions for better readability
2025-08-31 15:58:59 +08:00

108 lines
4.3 KiB
Python

import json
import re
from typing import List, Dict, Any
from flowllm import C, BaseLLMOp
from loguru import logger
from reme_ai.schema import Message
from reme_ai.schema.memory import BaseMemory
@C.register_op()
class MemoryValidationOp(BaseLLMOp):
file_path: str = __file__
def execute(self):
"""Validate quality of extracted task memories"""
task_memories: List[BaseMemory] = []
task_memories.extend(self.context.get("success_task_memories", []))
task_memories.extend(self.context.get("failure_task_memories", []))
task_memories.extend(self.context.get("comparative_task_memories", []))
if not task_memories:
logger.info("No task memories found for validation")
return
logger.info(f"Validating {len(task_memories)} extracted task memories")
# Validate task memories
validated_task_memories = []
for task_memory in task_memories:
validation_result = self._validate_single_task_memory(task_memory)
if validation_result and validation_result.get("is_valid", False):
task_memory.score = validation_result.get("score", 0.0)
validated_task_memories.append(task_memory)
else:
reason = validation_result.get("reason", "Unknown reason") if validation_result else "Validation failed"
logger.warning(f"Task memory validation failed: {reason}")
logger.info(f"Validated {len(validated_task_memories)} out of {len(task_memories)} task memories")
# Update context
self.context.response.answer = json.dumps([x.model_dump() for x in validated_task_memories])
self.context.response.metadata["memory_list"] = validated_task_memories
def _validate_single_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
"""Validate single task memory"""
validation_info = self._llm_validate_task_memory(task_memory)
logger.info(f"Validating: {validation_info}")
return validation_info
def _llm_validate_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
"""Validate task memory using LLM"""
try:
prompt = self.prompt_format(
prompt_name="task_memory_validation_prompt",
condition=task_memory.when_to_use,
task_memory_content=task_memory.content)
def parse_validation(message: Message) -> Dict[str, Any]:
try:
response_content = message.content
# Parse validation result
# Extract JSON blocks
json_pattern = r'```json\s*([\s\S]*?)\s*```'
json_blocks = re.findall(json_pattern, response_content)
if json_blocks:
parsed = json.loads(json_blocks[0])
else:
parsed = {}
is_valid = parsed.get("is_valid", True)
score = parsed.get("score", 0.5)
# Set validation threshold
validation_threshold = self.op_params.get("validation_threshold", 0.5)
return {
"is_valid": is_valid and score >= validation_threshold,
"score": score,
"feedback": response_content,
"reason": "" if (
is_valid and score >= validation_threshold) else f"Low validation score ({score:.2f}) or marked as invalid"
}
except Exception as e_inner:
logger.exception(f"Error parsing validation response: {e_inner}")
return {
"is_valid": False,
"score": 0.0,
"feedback": "",
"reason": f"Parse error: {str(e_inner)}"
}
return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_validation)
except Exception as e:
logger.error(f"LLM validation failed: {e}")
return {
"is_valid": False,
"score": 0.0,
"feedback": "",
"reason": f"LLM validation error: {str(e)}"
}