X-Agents/agent/app/agent/multi/workers/review.py

"""
Review Worker - 结果检查和质量评审
"""
import json
import re
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import HumanMessage, SystemMessage

from .base import BaseWorker
from ..types import AgentState, TaskItem, TaskStatus, ReviewResult
from ..prompts import REVIEW_SYSTEM_PROMPT


class ReviewWorker(BaseWorker):
    """Review Worker - 结果检查和质量评审"""

    def __init__(
        self,
        llm: BaseChatModel,
        tool_registry=None,
        tools: list = None
    ):
        super().__init__(
            llm=llm,
            name="review",
            system_prompt=REVIEW_SYSTEM_PROMPT,
            tools=tools or [],
            tool_registry=tool_registry
        )

    async def execute(self, task: TaskItem, context: dict) -> dict:
        """执行评审任务"""

        # 获取当前任务索引和任务计划
        # 注意：这里需要从 context 中获取更多信息

        # 构建 prompt
        context_str = json.dumps(context, ensure_ascii=False, indent=2) if context else "无"

        prompt = REVIEW_SYSTEM_PROMPT.format(
            original_task=context.get("original_task", ""),
            task_description=task.description,
            execution_result=task.result if task.result else "无结果",
            context=context_str
        )

        try:
            # 调用 LLM 进行评审
            response = await self.llm.ainvoke([
                SystemMessage(content=prompt),
                HumanMessage(content="请评审以上执行结果。")
            ])

            # 解析评审结果
            review_result = self._parse_review_response(response.content)

            # 根据评审结果决定下一步
            if review_result.passed:
                # 通过，更新任务状态为 completed
                new_status = TaskStatus.COMPLETED
                next_node = "supervisor"  # 返回 Supervisor 继续执行
            else:
                # 未通过，检查是否可重试
                if review_result.retryable:
                    new_status = TaskStatus.NEEDS_RETRY
                    next_node = "supervisor"  # 返回 Supervisor 决定是否重试
                else:
                    new_status = TaskStatus.FAILED
                    next_node = "aggregator"  # 失败，进入汇总

            return {
                "success": review_result.passed,
                "content": response.content,
                "review_result": review_result.model_dump() if hasattr(review_result, 'model_dump') else dict(review_result),
                "context": {
                    "review_passed": review_result.passed,
                    "issues": review_result.issues,
                    "last_review_by": self.name
                }
            }

        except Exception as e:
            return {
                "success": False,
                "content": "",
                "error": str(e),
                "context": {}
            }

    def _parse_review_response(self, response: str) -> ReviewResult:
        """解析评审响应"""
        try:
            # 尝试提取 JSON
            json_match = re.search(r'\{[\s\S]*\}', response)
            if json_match:
                data = json.loads(json_match.group())
            else:
                raise ValueError("No JSON found")

            return ReviewResult(
                passed=data.get("passed", True),
                issues=data.get("issues", []),
                suggestions=data.get("suggestions", []),
                retryable=data.get("retryable", True)
            )

        except Exception:
            # 解析失败，默认通过
            return ReviewResult(
                passed=True,
                issues=[],
                suggestions=[],
                retryable=True
            )

    def create_node(self):
        """创建 LangGraph 节点"""
        async def node(state: AgentState) -> dict:
            task_index = state.get("current_task_index", 0)
            task_plan = state.get("task_plan", [])

            if task_index >= len(task_plan):
                return {"next_node": "aggregator"}

            task = task_plan[task_index]
            shared_context = {
                **state.get("shared_context", {}),
                "original_task": state.get("original_task", "")
            }

            try:
                # 执行评审
                result = await self.execute(task, shared_context)

                # 更新任务状态
                review_passed = result.get("review_result", {}).get("passed", True)
                retryable = result.get("review_result", {}).get("retryable", True)

                if review_passed:
                    updated_status = TaskStatus.COMPLETED
                elif retryable:
                    updated_status = TaskStatus.NEEDS_RETRY
                else:
                    updated_status = TaskStatus.FAILED

                updated_plan = self._update_task_status(
                    task_plan,
                    task.id,
                    updated_status,
                    result=task.result
                )

                # 确定下一步
                if updated_status == TaskStatus.COMPLETED:
                    next_node = "supervisor"
                elif updated_status == TaskStatus.NEEDS_RETRY:
                    next_node = "supervisor"
                else:
                    next_node = "aggregator"

                return {
                    "task_plan": updated_plan,
                    "results": {**state.get("results", {}), f"{task.id}_review": result},
                    "shared_context": {**shared_context, **result.get("context", {})},
                    "next_node": next_node
                }

            except Exception as e:
                return {
                    "next_node": "aggregator",
                    "results": {**state.get("results", {}), f"{task.id}_review": {"error": str(e)}}
                }

        return node