diff --git a/src/agent/plan_execute/escalation.py b/src/agent/plan_execute/escalation.py new file mode 100644 index 00000000..be319060 --- /dev/null +++ b/src/agent/plan_execute/escalation.py @@ -0,0 +1,184 @@ +"""Deterministic escalation signal extraction for plan-execute runs.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Iterable + +from .models import Plan, StepResult + +DEFAULT_SPECIALIST_SERVERS = frozenset({"fmsr", "tsfm", "vibration", "wo"}) + +DEFAULT_ESCALATION_TERMS = ( + "work order", + "work-order", + "diagnostic", + "diagnostics", + "diagnosis", + "failure", + "failed", + "fault", + "alarm", + "anomaly", +) + + +@dataclass +class EscalationSignals: + """Signals that may later inform adaptive escalation decisions.""" + + step_count: int + dependency_depth: int + uses_specialist_servers: bool + specialist_servers_used: list[str] = field(default_factory=list) + any_step_failed: bool = False + failed_steps: list[int] = field(default_factory=list) + servers_used: list[str] = field(default_factory=list) + tools_used: list[str] = field(default_factory=list) + matched_terms: list[str] = field(default_factory=list) + + @property + def has_domain_terms(self) -> bool: + return bool(self.matched_terms) + + +@dataclass +class EscalationDecision: + """Deterministic escalation decision for a plan-execute run.""" + + should_escalate: bool + reasons: list[str] = field(default_factory=list) + + +def extract_escalation_signals( + question: str, + plan: Plan, + trajectory: Iterable[StepResult] | None = None, + specialist_servers: Iterable[str] = DEFAULT_SPECIALIST_SERVERS, + escalation_terms: Iterable[str] = DEFAULT_ESCALATION_TERMS, +) -> EscalationSignals: + """Extract deterministic escalation signals from a plan and trajectory. + + This function is intentionally side-effect free and makes no LLM calls. + """ + results = list(trajectory or []) + specialist_server_set = set(specialist_servers) + + servers_used = _unique_sorted( + [step.server for step in plan.steps] + [result.server for result in results] + ) + tools_used = _unique_sorted( + [step.tool for step in plan.steps if step.tool] + + [result.tool for result in results if result.tool] + ) + specialist_servers_used = [ + server for server in servers_used if server in specialist_server_set + ] + failed_steps = [result.step_number for result in results if not result.success] + + return EscalationSignals( + step_count=len(plan.steps), + dependency_depth=_dependency_depth(plan), + uses_specialist_servers=bool(specialist_servers_used), + specialist_servers_used=specialist_servers_used, + any_step_failed=bool(failed_steps), + failed_steps=failed_steps, + servers_used=servers_used, + tools_used=tools_used, + matched_terms=_matched_terms(question, plan, results, escalation_terms), + ) + + +def should_escalate(signals: EscalationSignals) -> EscalationDecision: + """Return a deterministic adaptive-escalation decision for extracted signals.""" + reasons = [] + if signals.any_step_failed: + reasons.append("failed step") + if signals.dependency_depth >= 3: + reasons.append("dependency depth >= 3") + if signals.uses_specialist_servers: + reasons.append("specialist server used") + if signals.has_domain_terms: + reasons.append("domain escalation term matched") + + return EscalationDecision(should_escalate=bool(reasons), reasons=reasons) + + +def _dependency_depth(plan: Plan) -> int: + """Return the longest dependency chain length, counting the step itself.""" + steps_by_number = {step.step_number: step for step in plan.steps} + visiting: set[int] = set() + memo: dict[int, int] = {} + + def depth(step_number: int) -> int: + if step_number in memo: + return memo[step_number] + if step_number in visiting: + return 1 + + step = steps_by_number.get(step_number) + if step is None: + return 0 + + visiting.add(step_number) + dep_depth = max((depth(dep) for dep in step.dependencies), default=0) + visiting.remove(step_number) + memo[step_number] = dep_depth + 1 + return memo[step_number] + + return max((depth(step.step_number) for step in plan.steps), default=0) + + +def _matched_terms( + question: str, + plan: Plan, + trajectory: list[StepResult], + escalation_terms: Iterable[str], +) -> list[str]: + text = "\n".join( + [ + question, + plan.raw, + *[ + "\n".join([step.task, step.expected_output, step.server, step.tool]) + for step in plan.steps + ], + *[ + "\n".join( + [ + result.task, + result.server, + result.tool, + result.response, + result.error or "", + ] + ) + for result in trajectory + ], + ] + ) + matched = [] + seen = set() + for term in escalation_terms: + key = term.casefold() + if key in seen: + continue + if re.search(rf"(? list[str]: + return sorted({value for value in values if value}) + + +__all__ = [ + "DEFAULT_ESCALATION_TERMS", + "DEFAULT_SPECIALIST_SERVERS", + "EscalationDecision", + "EscalationSignals", + "extract_escalation_signals", + "should_escalate", +] diff --git a/src/agent/plan_execute/runner.py b/src/agent/plan_execute/runner.py index 445f47b0..ac41af79 100644 --- a/src/agent/plan_execute/runner.py +++ b/src/agent/plan_execute/runner.py @@ -19,6 +19,7 @@ from llm import LLMBackend, LLMResult from observability import agent_run_span, persist_trajectory +from .escalation import extract_escalation_signals, should_escalate from .executor import Executor from .models import OrchestratorResult from .planner import Planner @@ -76,6 +77,40 @@ def model_id(self) -> str: above. Do not repeat the individual steps — just give the final answer. """ +_VERIFY_PROMPT = """\ +You are reviewing the evidence gathered by a plan-execute industrial asset \ +operations agent before it gives a final answer. + +Original question: {question} + +Escalation reasons: +{reasons} + +Step-by-step execution results: +{results} + +Check for failed steps, missing evidence, conflicting evidence, alternative \ +explanations, and whether any work-order or action recommendation would be \ +premature. Provide concise verification notes only. +""" + +_SUMMARIZE_WITH_VERIFICATION_PROMPT = """\ +You are summarizing the results of a multi-step task execution for an \ +industrial asset operations system. + +Original question: {question} + +Step-by-step execution results: +{results} + +Verification notes: +{verification} + +Provide a concise, direct answer to the original question based on the results \ +and verification notes above. Do not repeat the individual steps — just give \ +the final answer. +""" + class PlanExecuteRunner(AgentRunner): """Entry-point for plan-and-execute workflows using MCP servers as tool providers. @@ -95,14 +130,19 @@ class PlanExecuteRunner(AgentRunner): names the planner will assign steps to. Values are either a uv entry-point name (str) or a Path to a script file. Defaults to all five registered servers. + adaptive_escalation: Enable an experimental deterministic escalation + policy that can add a verification pass before + summarisation. Defaults to ``False``. """ def __init__( self, llm: LLMBackend, server_paths: dict[str, Path | str] | None = None, + adaptive_escalation: bool = False, ) -> None: super().__init__(llm, server_paths) + self._adaptive_escalation = adaptive_escalation self._meter = _TokenMeter(llm) self._planner = Planner(self._meter) self._executor = Executor(self._meter, server_paths) @@ -143,6 +183,31 @@ async def run(self, question: str) -> OrchestratorResult: # 3. Execute trajectory = await self._executor.execute_plan(plan, question) + span.set_attribute("agent.escalation.enabled", self._adaptive_escalation) + verification = "" + escalation_decision = None + if self._adaptive_escalation: + signals = extract_escalation_signals(question, plan, trajectory) + escalation_decision = should_escalate(signals) + span.set_attribute( + "agent.escalation.should_escalate", + escalation_decision.should_escalate, + ) + span.set_attribute( + "agent.escalation.reasons", + escalation_decision.reasons, + ) + span.set_attribute( + "agent.escalation.dependency_depth", signals.dependency_depth + ) + span.set_attribute( + "agent.escalation.any_step_failed", signals.any_step_failed + ) + span.set_attribute( + "agent.escalation.uses_specialist_servers", + signals.uses_specialist_servers, + ) + # 4. Summarise _log.info("Summarising...") results_text = "\n\n".join( @@ -151,9 +216,26 @@ async def run(self, question: str) -> OrchestratorResult: for r in trajectory ) summarization_started = time.perf_counter() - answer = self._meter.generate( - _SUMMARIZE_PROMPT.format(question=question, results=results_text) - ) + if escalation_decision and escalation_decision.should_escalate: + _log.info("Running adaptive escalation verification...") + verification = self._meter.generate( + _VERIFY_PROMPT.format( + question=question, + reasons="\n".join(escalation_decision.reasons), + results=results_text, + ) + ) + answer = self._meter.generate( + _SUMMARIZE_WITH_VERIFICATION_PROMPT.format( + question=question, + results=results_text, + verification=verification, + ) + ) + else: + answer = self._meter.generate( + _SUMMARIZE_PROMPT.format(question=question, results=results_text) + ) summarization_ms = (time.perf_counter() - summarization_started) * 1000 duration_ms = (time.perf_counter() - run_started) * 1000 diff --git a/src/agent/tests/test_escalation.py b/src/agent/tests/test_escalation.py new file mode 100644 index 00000000..7c034f04 --- /dev/null +++ b/src/agent/tests/test_escalation.py @@ -0,0 +1,234 @@ +"""Tests for deterministic escalation signal extraction.""" + +from agent.plan_execute.escalation import ( + DEFAULT_SPECIALIST_SERVERS, + EscalationDecision, + extract_escalation_signals, + should_escalate, +) +from agent.plan_execute.models import Plan, PlanStep, StepResult +from agent.runner import DEFAULT_SERVER_PATHS + + +def _step( + n: int, + server: str = "iot", + tool: str = "sites", + deps: list[int] | None = None, + task: str | None = None, + expected_output: str = "output", +) -> PlanStep: + return PlanStep( + step_number=n, + task=task or f"Task {n}", + server=server, + tool=tool, + tool_args={}, + dependencies=deps or [], + expected_output=expected_output, + ) + + +def test_extracts_step_count_and_dependency_depth(): + plan = Plan( + steps=[ + _step(1), + _step(2, deps=[1]), + _step(3, deps=[2]), + _step(4, deps=[1]), + ], + raw="", + ) + + signals = extract_escalation_signals("Q", plan) + + assert signals.step_count == 4 + assert signals.dependency_depth == 3 + + +def test_empty_plan_has_zero_dependency_depth(): + signals = extract_escalation_signals("Q", Plan(steps=[], raw="")) + + assert signals.step_count == 0 + assert signals.dependency_depth == 0 + + +def test_detects_specialist_servers(): + plan = Plan( + steps=[ + _step(1, server="iot"), + _step(2, server="wo", tool="work_orders"), + _step(3, server="vibration", tool="analyze"), + ], + raw="", + ) + + signals = extract_escalation_signals("Q", plan) + + assert signals.uses_specialist_servers is True + assert signals.specialist_servers_used == ["vibration", "wo"] + + +def test_default_specialist_servers_match_registered_server_names(): + assert DEFAULT_SPECIALIST_SERVERS <= set(DEFAULT_SERVER_PATHS) + + +def test_collects_servers_and_tools_from_plan_and_trajectory(): + plan = Plan( + steps=[ + _step(1, server="iot", tool="assets"), + _step(2, server="utilities", tool="current_date_time"), + ], + raw="", + ) + trajectory = [ + StepResult( + step_number=1, + task="Task 1", + server="iot", + response="ok", + tool="assets", + ), + StepResult( + step_number=2, + task="Task 2", + server="fmsr", + response="ok", + tool="diagnose_failure", + ), + ] + + signals = extract_escalation_signals("Q", plan, trajectory) + + assert signals.servers_used == ["fmsr", "iot", "utilities"] + assert signals.tools_used == ["assets", "current_date_time", "diagnose_failure"] + + +def test_detects_failed_steps_from_trajectory(): + plan = Plan(steps=[_step(1), _step(2)], raw="") + trajectory = [ + StepResult(step_number=1, task="Task 1", server="iot", response="ok"), + StepResult( + step_number=2, + task="Task 2", + server="iot", + response="", + error="timeout", + ), + ] + + signals = extract_escalation_signals("Q", plan, trajectory) + + assert signals.any_step_failed is True + assert signals.failed_steps == [2] + + +def test_matches_domain_terms_across_question_plan_and_trajectory(): + plan = Plan( + steps=[ + _step( + 1, + task="Open work order history", + expected_output="Recent maintenance records", + ) + ], + raw="#Task1: Run diagnostics", + ) + trajectory = [ + StepResult( + step_number=1, + task="Task 1", + server="iot", + response="Asset reported a pump failure alarm", + ) + ] + + signals = extract_escalation_signals("Any anomaly on CH-1?", plan, trajectory) + + assert signals.has_domain_terms is True + assert signals.matched_terms == [ + "work order", + "diagnostics", + "failure", + "alarm", + "anomaly", + ] + + +def test_custom_specialist_servers_and_terms_are_supported(): + plan = Plan(steps=[_step(1, server="custom", task="Check severe drift")], raw="") + + signals = extract_escalation_signals( + "Q", + plan, + specialist_servers={"custom"}, + escalation_terms=["severe drift"], + ) + + assert signals.uses_specialist_servers is True + assert signals.specialist_servers_used == ["custom"] + assert signals.matched_terms == ["severe drift"] + + +def test_matched_terms_are_case_insensitive_and_deduplicated(): + plan = Plan(steps=[_step(1, task="Investigate FAILURE alarm")], raw="") + + signals = extract_escalation_signals( + "Failure reported", + plan, + escalation_terms=["failure", "Failure", "alarm"], + ) + + assert signals.matched_terms == ["failure", "alarm"] + + +def test_escalation_decision_dataclass_is_available(): + decision = EscalationDecision(should_escalate=True, reasons=["failed step"]) + + assert decision.should_escalate is True + assert decision.reasons == ["failed step"] + + +def test_policy_escalates_on_failed_steps(): + plan = Plan(steps=[_step(1)], raw="") + trajectory = [ + StepResult( + step_number=1, + task="Task 1", + server="iot", + response="", + error="timeout", + ) + ] + + decision = should_escalate(extract_escalation_signals("Q", plan, trajectory)) + + assert decision.should_escalate is True + assert decision.reasons == ["failed step"] + + +def test_policy_escalates_on_specialist_server_usage(): + plan = Plan(steps=[_step(1, server="vibration")], raw="") + + decision = should_escalate(extract_escalation_signals("Q", plan)) + + assert decision.should_escalate is True + assert decision.reasons == ["specialist server used"] + + +def test_policy_escalates_on_domain_terms(): + plan = Plan(steps=[_step(1, task="Review work order history")], raw="") + + decision = should_escalate(extract_escalation_signals("Q", plan)) + + assert decision.should_escalate is True + assert decision.reasons == ["domain escalation term matched"] + + +def test_policy_does_not_escalate_simple_low_risk_plan(): + plan = Plan(steps=[_step(1, task="List sites", expected_output="Site list")], raw="") + + decision = should_escalate(extract_escalation_signals("Q", plan)) + + assert decision.should_escalate is False + assert decision.reasons == [] diff --git a/src/agent/tests/test_runner.py b/src/agent/tests/test_runner.py index 062d7865..2ad98ab8 100644 --- a/src/agent/tests/test_runner.py +++ b/src/agent/tests/test_runner.py @@ -96,6 +96,18 @@ def generate(self, prompt: str, **_kw) -> str: return self._response +class _RecordingSequentialLLM(LLMBackend): + """Records prompts while returning canned responses in order.""" + + def __init__(self, responses: list[str]) -> None: + self.prompts: list[str] = [] + self._responses = iter(responses) + + def generate(self, prompt: str, temperature: float = 0.0) -> str: + self.prompts.append(prompt) + return next(self._responses, "") + + # ── orchestrator tests ──────────────────────────────────────────────────────── @@ -146,6 +158,54 @@ async def test_orchestrator_unknown_server_recorded_as_error(sequential_llm): assert "ghost" in result.trajectory[0].error +def test_orchestrator_adaptive_escalation_disabled_by_default(mock_llm): + runner = PlanExecuteRunner(mock_llm()) + + assert runner._adaptive_escalation is False + + +@pytest.mark.anyio +async def test_orchestrator_adaptive_escalation_no_extra_call_for_low_risk_plan(): + plan = ( + "#Task1: List sites\n" + "#Server1: iot\n" + "#Tool1: sites\n" + "#Dependency1: None\n" + "#ExpectedOutput1: Site list\n" + ) + llm = _RecordingSequentialLLM([plan, "{}", _FINAL_ANSWER]) + + with _patch_mcp()[0], _patch_mcp()[1]: + result = await PlanExecuteRunner(llm, adaptive_escalation=True).run("Q") + + assert result.answer == _FINAL_ANSWER + assert len(llm.prompts) == 3 + assert not any("Verification notes" in prompt for prompt in llm.prompts) + + +@pytest.mark.anyio +async def test_orchestrator_adaptive_escalation_runs_verification_when_enabled(): + plan = ( + "#Task1: Check vibration trend\n" + "#Server1: vibration\n" + "#Tool1: analyze\n" + "#Dependency1: None\n" + "#ExpectedOutput1: Vibration analysis\n" + ) + verification = "Verification: evidence is limited." + final_answer = "Final answer with verification." + llm = _RecordingSequentialLLM([plan, "{}", verification, final_answer]) + + with _patch_mcp()[0], _patch_mcp()[1]: + result = await PlanExecuteRunner(llm, adaptive_escalation=True).run("Q") + + assert result.answer == final_answer + assert len(llm.prompts) == 4 + assert "Escalation reasons:" in llm.prompts[2] + assert "specialist server used" in llm.prompts[2] + assert verification in llm.prompts[3] + + class _UsageReportingLLM(LLMBackend): """Sequential LLM that reports per-call token usage via LLMResult."""