agent-platform/core/engine/nodes.py

"""
LangGraph node functions for agent execution flow.

Each node is decorated with @agent_span for OpenTelemetry tracing.
All nodes follow the error boundary pattern: catch all exceptions,
set state.error, and return state (never let exceptions escape).
"""

import json
import logging
from decimal import Decimal
from typing import Any

from langgraph.errors import GraphInterrupt

from core.engine.state import AgentState

# Try to import OTel decorator, fallback to passthrough if not available
try:
    from core.observability.tracer import agent_span
except ImportError:
    from collections.abc import Callable
    from typing import TypeVar

    F = TypeVar('F', bound=Callable[..., Any])

    def agent_span(node_name: str) -> Callable[[F], F]:
        def decorator(func: F) -> F:
            return func
        return decorator

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════
# Tool Definitions (inline as per spec)
# ═══════════════════════════════════════════════════════════════

EXTRACT_PARAMS_TOOL = {
    "name": "extract_user_params",
    "description": (
        "从用户自然语言输入提取结构化业务参数。"
        "识别所有可提取的字段，标记缺失的必填字段，"
        "并为每个提取的值提供置信度评分。"
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "profile_type": {
                "type": "string",
                "enum": ["presales", "security_analysis", "knowledge_qa", "custom"],
                "description": "识别的业务场景类型",
            },
            "extracted_fields": {
                "type": "array",
                "description": "成功提取的字段列表",
                "items": {
                    "type": "object",
                    "properties": {
                        "field_name": {"type": "string", "description": "字段名称"},
                        "value": {"description": "提取的值（任意类型）"},
                        "confidence": {
                            "type": "number",
                            "minimum": 0,
                            "maximum": 1,
                            "description": "提取置信度 (0-1)",
                        },
                        "source_text": {
                            "type": "string",
                            "description": "原文中的相关文本片段",
                        },
                    },
                    "required": ["field_name", "value", "confidence"],
                },
            },
            "missing_required_fields": {
                "type": "array",
                "items": {"type": "string"},
                "description": "缺失的必填字段列表",
            },
            "task_summary": {
                "type": "string",
                "description": "用户任务的简短摘要",
            },
        },
        "required": ["profile_type", "extracted_fields", "task_summary"],
    },
}

GENERATE_PLAN_TOOL = {
    "name": "generate_execution_plan",
    "description": (
        "基于 Profile SOP 和提取的参数生成逐步执行计划。"
        "每个步骤指定工具名称、参数、依赖关系和审批要求。"
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "plan_id": {"type": "string", "description": "计划唯一标识符"},
            "steps": {
                "type": "array",
                "description": "执行步骤列表",
                "items": {
                    "type": "object",
                    "properties": {
                        "step_id": {"type": "string", "description": "步骤唯一标识符"},
                        "tool_name": {"type": "string", "description": "要调用的工具名称"},
                        "status": {
                            "type": "string",
                            "enum": ["pending", "skipped", "requires_approval"],
                            "description": "步骤初始状态",
                        },
                        "params": {
                            "type": "object",
                            "description": "工具参数（支持模板语法）",
                        },
                        "output_key": {
                            "type": "string",
                            "description": "结果存储键名",
                        },
                        "on_failure": {
                            "type": "string",
                            "enum": ["abort", "skip"],
                            "description": "失败策略",
                        },
                        "depends_on": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "依赖的步骤 ID 列表",
                        },
                        "approval_required": {
                            "type": "boolean",
                            "description": "是否需要人工审批",
                        },
                        "estimated_duration_ms": {
                            "type": "integer",
                            "description": "预估执行时长（毫秒）",
                        },
                    },
                    "required": ["step_id", "tool_name", "status", "params", "output_key"],
                },
            },
            "plan_summary": {"type": "string", "description": "计划整体摘要"},
        },
        "required": ["plan_id", "steps", "plan_summary"],
    },
}


# ═══════════════════════════════════════════════════════════════
# Helper Functions
# ═══════════════════════════════════════════════════════════════


def _resolve_params(params: dict[str, Any], state: AgentState) -> dict[str, Any]:
    """
    Resolve template syntax in tool parameters.

    Supported patterns:
    - {{user_input}} → state.user_input
    - {{user_input.field}} → state.extracted_params.get("field")
    - {{steps.step_id.field}} → state.step_outputs["step_id"].get("field")
    - {{profile.field}} → (logged, kept as-is)

    Unresolvable keys are kept as-is (no exception raised).
    """
    import re

    unresolved: list[str] = []

    def resolve_value(value: Any) -> Any:
        if isinstance(value, str):
            pattern = r"\{\{([^}]+)\}\}"
            matches = re.findall(pattern, value)

            for match in matches:
                parts = match.strip().split(".")
                resolved = None

                if parts[0] == "user_input":
                    if len(parts) == 1:
                        resolved = state.user_input
                    elif len(parts) == 2:
                        resolved = state.extracted_params.get(parts[1])
                elif parts[0] == "steps" and len(parts) >= 3:
                    step_id = parts[1]
                    field = parts[2]
                    step_data = state.step_outputs.get(step_id, {})
                    resolved = step_data.get(field)
                elif parts[0] == "profile":
                    logger.warning("Profile template not resolved: %s", match)
                    unresolved.append(match)
                    continue

                if resolved is not None:
                    value = value.replace(f"{{{{{match}}}}}", str(resolved))
                else:
                    logger.warning("Unresolved template: %s", match)
                    unresolved.append(match)

            return value
        elif isinstance(value, dict):
            return {k: resolve_value(v) for k, v in value.items()}
        elif isinstance(value, list):
            return [resolve_value(item) for item in value]
        else:
            return value

    resolved_params = resolve_value(params)

    if unresolved and isinstance(resolved_params, dict):
        if "meta" not in resolved_params:
            resolved_params["meta"] = {}
        resolved_params["meta"]["unresolved_templates"] = unresolved

    return resolved_params  # type: ignore[no-any-return]


# ═══════════════════════════════════════════════════════════════
# Node Functions
# ═══════════════════════════════════════════════════════════════


@agent_span("ask")
async def ask_node(state: AgentState) -> AgentState:
    """
    Extract structured parameters from user input.

    Calls Claude with EXTRACT_PARAMS_TOOL to identify:
    - Extracted fields with confidence scores
    - Missing required fields

    If required fields are missing, sets output_type to "clarification_needed"
    and routes to END.  Otherwise, populates state.extracted_params.
    """
    try:
        from core.llm.claude_client import ClaudeClient
        from core.llm.prompt_builder import build_ask_user_message, build_system_messages
        from core.profile.registry import ProfileRegistry

        # Explicitly initialize optional fields so LangGraph tracks them in all code paths.
        state.error = None
        state.pending_approval = None

        logger.info("[ask_node] trace_id=%s user_input=%.100s", state.trace_id, state.user_input)

        profile = ProfileRegistry.instance().get(state.profile_type)

        system_msgs = build_system_messages(profile.model_dump())
        schema = {"input_fields": profile.input_fields} if profile.input_fields else {}
        user_msg = build_ask_user_message(state.user_input, schema)

        client = ClaudeClient()
        tool_input, cost = await client.call_with_tools(
            system_msgs=system_msgs,
            user_msg=user_msg,
            tools=[EXTRACT_PARAMS_TOOL],
            tool_choice={"type": "any"},
            trace_id=state.trace_id,
        )

        extracted_fields = tool_input.get("extracted_fields", [])
        missing_fields = tool_input.get("missing_required_fields", [])

        if missing_fields:
            state.output_type = "clarification_needed"
            missing_list = ", ".join(missing_fields)
            state.final_output = (
                f"需要补充以下信息才能继续：{missing_list}。请提供这些字段的值。"
            )
            logger.info("[ask_node] Clarification needed: %s", missing_fields)
            return state

        state.extracted_params = {
            field["field_name"]: field["value"] for field in extracted_fields
        }
        state.total_cost_cny += cost.cost_cny

        logger.info("[ask_node] Extracted params: %s", list(state.extracted_params.keys()))
        return state

    except Exception as e:
        logger.exception("[ask_node] Unhandled exception: %s", e)
        state.error = {
            "code": "ASK_NODE_ERROR",
            "message": str(e),
            "node": "ask",
            "trace_id": state.trace_id,
        }
        return state


@agent_span("plan")
async def plan_node(state: AgentState) -> AgentState:
    """
    Generate execution plan based on Profile SOP.

    Calls Claude with GENERATE_PLAN_TOOL to create step-by-step plan.
    Tool whitelist enforcement is delegated to ToolExecutor at execution time.
    """
    try:
        from core.llm.claude_client import ClaudeClient
        from core.llm.prompt_builder import build_plan_message, build_system_messages
        from core.profile.registry import ProfileRegistry

        logger.info("[plan_node] trace_id=%s user_role=%s", state.trace_id, state.user_role)

        profile = ProfileRegistry.instance().get(state.profile_type)

        # Build SOP text from profile description + workflow steps
        sop_text = profile.description + "\n\n工作流步骤:\n" + json.dumps(
            [s.model_dump() for s in profile.workflow_steps],
            ensure_ascii=False,
            indent=2,
        )

        system_msgs = build_system_messages(profile.model_dump())
        user_msg = build_plan_message(
            extracted_params=state.extracted_params,
            sop=sop_text,
            user_permissions={"role": state.user_role, "data_scope": state.data_scope},
        )

        client = ClaudeClient()
        tool_input, cost = await client.call_with_tools(
            system_msgs=system_msgs,
            user_msg=user_msg,
            tools=[GENERATE_PLAN_TOOL],
            tool_choice={"type": "any"},
            trace_id=state.trace_id,
        )

        steps = tool_input.get("steps", [])
        if not steps:
            state.error = {
                "code": "PLAN_EMPTY",
                "message": "生成的执行计划为空",
                "node": "plan",
                "trace_id": state.trace_id,
            }
            return state

        state.execution_plan = steps
        state.total_cost_cny += cost.cost_cny

        logger.info("[plan_node] Generated %d steps", len(steps))
        return state

    except Exception as e:
        logger.exception("[plan_node] Unhandled exception: %s", e)
        state.error = {
            "code": "PLAN_NODE_ERROR",
            "message": str(e),
            "node": "plan",
            "trace_id": state.trace_id,
        }
        return state


@agent_span("act")
async def act_node(state: AgentState) -> AgentState:
    """
    Execute tools according to the execution plan.

    Processes steps sequentially from current_step_index:
    1. Resolve parameter templates
    2. Check for approval requirements (interrupt() if needed)
    3. Execute tool via ToolExecutor
    4. Store result summary + citations in step_outputs
    5. Handle failures based on on_failure strategy
    """
    try:
        from langgraph.types import interrupt

        from core.audit.logger import AuditLogger
        from core.profile.registry import ProfileRegistry
        from core.tools.base import ToolNotAllowedError
        from core.tools.executor import ToolExecutor
        from core.tools.registry import get_registry

        logger.info(
            "[act_node] trace_id=%s step_index=%d",
            state.trace_id,
            state.current_step_index,
        )

        profile = ProfileRegistry.instance().get(state.profile_type)
        tool_registry = get_registry()
        audit_logger = AuditLogger()
        tool_executor = ToolExecutor(registry=tool_registry, audit_logger=audit_logger)

        while state.current_step_index < len(state.execution_plan):
            step = state.execution_plan[state.current_step_index]
            step_id = step["step_id"]
            tool_name = step["tool_name"]

            logger.info("[act_node] Executing step %s: %s", step_id, tool_name)

            resolved_params = _resolve_params(step.get("params", {}), state)

            # ── Approval gate ──────────────────────────────────────────────────
            if step.get("approval_required", False):
                approval_payload = {
                    "step_id": step_id,
                    "tool_name": tool_name,
                    "params": resolved_params,
                    "reason": f"步骤 {step_id} 需要人工审批",
                }
                state.pending_approval = approval_payload
                logger.info("[act_node] Approval required for step %s, interrupting", step_id)
                # Pause until resumed with Command(resume=True/False)
                approved = interrupt(approval_payload)
                if not approved:
                    state.error = {
                        "code": "APPROVAL_REJECTED",
                        "message": f"步骤 {step_id} 审批被拒绝",
                        "node": "act",
                        "trace_id": state.trace_id,
                    }
                    return state
                state.pending_approval = None

            # ── Tool execution ─────────────────────────────────────────────────
            try:
                result = await tool_executor.execute_tool(
                    tool_name=tool_name,
                    raw_params=resolved_params,
                    allowed_tools=profile.allowed_tools,
                    user_id=state.user_id,
                    tenant_id=state.tenant_id,
                    data_scope=state.data_scope[0] if state.data_scope else "",
                    trace_id=state.trace_id,
                    step_id=step_id,
                )
            except ToolNotAllowedError as e:
                state.error = {
                    "code": "TOOL_NOT_ALLOWED",
                    "message": str(e),
                    "tool": tool_name,
                    "node": "act",
                    "trace_id": state.trace_id,
                }
                logger.error("[act_node] Tool not allowed: %s", tool_name)
                return state

            # ── Audit every tool call (success or failure) ──────────────────────
            await audit_logger.write_tool_call(
                trace_id=state.trace_id,
                step_id=step_id,
                tool_name=tool_name,
                result=result,
                actor=state.user_id,
                tenant_id=state.tenant_id,
            )

            if not result.success:
                on_failure = step.get("on_failure", "abort")
                if on_failure == "abort":
                    state.error = {
                        "code": "TOOL_FAILED",
                        "message": result.error or "工具执行失败",
                        "tool": tool_name,
                        "step_id": step_id,
                        "node": "act",
                        "trace_id": state.trace_id,
                    }
                    logger.error("[act_node] Tool %s failed, aborting", tool_name)
                    return state
                else:
                    logger.warning(
                        "[act_node] Tool %s failed, continuing (on_failure=%s)",
                        tool_name,
                        on_failure,
                    )

            # ── Store result ───────────────────────────────────────────────────
            output_key = step["output_key"]
            summary = (
                result.data.get("summary", str(result.data)[:200])
                if isinstance(result.data, dict)
                else str(result.data)[:200] if result.data
                else "执行完成"
            )
            data_ref = result.data.get("data_ref") if isinstance(result.data, dict) else None
            state.add_step_output(output_key, summary, data_ref)

            # Store citations for evidence scoring (mode="json" converts Decimal → str)
            if result.citations:
                state.step_outputs[output_key]["citations"] = [
                    c.model_dump(mode="json") for c in result.citations
                ]

            state.current_step_index += 1

        logger.info("[act_node] All steps completed")
        return state

    except GraphInterrupt:
        # LangGraph interrupt must propagate — do not catch as an error.
        raise
    except Exception as e:
        logger.exception("[act_node] Unhandled exception: %s", e)
        state.error = {
            "code": "ACT_NODE_ERROR",
            "message": str(e),
            "node": "act",
            "trace_id": state.trace_id,
        }
        return state


@agent_span("verify")
async def verify_node(state: AgentState) -> AgentState:
    """
    Assess evidence quality of execution results.

    Collects citations from all step outputs and scores them via EvidenceScorer.
    Sets evidence_score, evidence_level, and verified_citations.
    """
    try:
        from core.profile.registry import ProfileRegistry
        from core.verify.scorer import EvidenceScorer

        logger.info("[verify_node] trace_id=%s", state.trace_id)

        profile = ProfileRegistry.instance().get(state.profile_type)

        # Collect all citations from step outputs
        all_citations: list[dict[str, Any]] = []
        for step_output in state.step_outputs.values():
            if isinstance(step_output, dict) and "citations" in step_output:
                all_citations.extend(step_output["citations"])

        scorer = EvidenceScorer()
        scorer_result = scorer.score(
            citations=all_citations,
            step_outputs=state.step_outputs,
            required_fields=profile.required_output_fields,
            historical_accuracy=Decimal("0.8"),
        )

        state.evidence_score = scorer_result.score
        state.evidence_level = scorer_result.level
        state.citations = [c.model_dump() for c in scorer_result.verified_citations]

        logger.info(
            "[verify_node] Evidence score=%s level=%s",
            state.evidence_score,
            state.evidence_level,
        )
        return state

    except Exception as e:
        logger.exception("[verify_node] Unhandled exception: %s", e)
        state.error = {
            "code": "VERIFY_NODE_ERROR",
            "message": str(e),
            "node": "verify",
            "trace_id": state.trace_id,
        }
        return state


@agent_span("report")
async def report_node(state: AgentState) -> AgentState:
    """
    Generate final output based on evidence level.

    If evidence is insufficient, sets output_type to "insufficient_data"
    and explains what's missing.  Otherwise, calls Claude to generate
    a formatted report using the output template.
    """
    try:
        from core.llm.claude_client import ClaudeClient
        from core.llm.prompt_builder import build_report_message, build_system_messages
        from core.profile.registry import ProfileRegistry

        logger.info(
            "[report_node] trace_id=%s evidence_level=%s",
            state.trace_id,
            state.evidence_level,
        )

        profile = ProfileRegistry.instance().get(state.profile_type)

        if state.evidence_level == "insufficient":
            state.output_type = "insufficient_data"
            state.final_output = (
                "证据不足，无法生成可靠的输出。\n"
                "缺失的数据：请补充更多相关信息或调整查询条件。"
            )
            logger.info("[report_node] Insufficient evidence, cannot generate report")
            return state

        system_msgs = build_system_messages(profile.model_dump())
        user_msg = build_report_message(
            evidence_level=state.evidence_level or "sufficient",
            verified_data=state.step_outputs,
            citations=state.citations,
            template_hint=profile.output_template,
        )

        client = ClaudeClient()
        text, cost = await client.call_text(
            system_msgs=system_msgs,
            user_msg=user_msg,
            max_tokens=4096,
            trace_id=state.trace_id,
        )

        state.final_output = text
        state.output_type = "document" if profile.output_format == "docx" else "text"
        state.total_cost_cny += cost.cost_cny

        logger.info("[report_node] Report generated, length=%d", len(state.final_output))
        return state

    except Exception as e:
        logger.exception("[report_node] Unhandled exception: %s", e)
        state.error = {
            "code": "REPORT_NODE_ERROR",
            "message": str(e),
            "node": "report",
            "trace_id": state.trace_id,
        }
        return state