backend/agents/executor.py

"""
backend/agents/executor.py

The Executor Agent — executes one plan step at a time.

Responsibilities:
  1. Take the current pending step
  2. Decide HOW to execute it (which tool, what args)
  3. Call the tool via function calling
  4. Handle errors with retry logic
  5. Store result in state
  6. Mark step complete or failed
"""
from __future__ import annotations
import asyncio
import json
import time
from datetime import datetime, timezone

from langchain_core.messages import SystemMessage, HumanMessage

from ..core.llm import get_llm
from ..state.graph_state import (
    WorkflowState, TaskStatus, StepStatus, AgentRole, make_agent_event
)
from ..tools.registry import TOOL_SCHEMAS, execute_tool
from ..core.config import get_settings
from ..core.logger import get_logger

log = get_logger(__name__)

EXECUTOR_SYSTEM = """Execute ONE step using the available tools. Choose the right tool and args. Always call a tool."""

SYNTHESIZE_SYSTEM = """Write the final answer to the task using the collected results. Be concise and clear.
Use plain text only - no LaTeX, no $\\boxed{}, no math notation. Write numbers as regular text (e.g. "1040" not "$\\boxed{1040}$")."""


def _build_context(state: WorkflowState) -> str:
    """Build context string from previous step results."""
    parts = [f"Original task: {state['task']}"]

    if state["step_results"]:
        parts.append("\nPrevious step results:")
        for step_id, result in state["step_results"].items():
            result_str = json.dumps(result, default=str)[:600]
            parts.append(f"  [{step_id}]: {result_str}")

    return "\n".join(parts)


def _get_current_step(state: WorkflowState) -> dict | None:
    """Find the next pending step that has all dependencies satisfied."""
    completed_ids = {s["step_id"] for s in state["plan"] if s["status"] == "done"}

    for step in state["plan"]:
        if step["status"] != "pending":
            continue
        deps_satisfied = all(d in completed_ids for d in step.get("depends_on", []))
        if deps_satisfied:
            return step
    return None


def executor_node(state: WorkflowState) -> WorkflowState:
    """LangGraph node — executes one pending plan step."""
    settings = get_settings()

    step = _get_current_step(state)
    if step is None:
        log.info("No pending steps", task_id=state["task_id"])
        return {**state, "status": TaskStatus.REFLECTING if settings.enable_reflection else TaskStatus.COMPLETED}

    log.info("Executor running step", step_id=step["step_id"], tool=step.get("tool"), title=step["title"])

    # Handle synthesize step specially
    if step.get("tool") == "synthesize":
        return _handle_synthesize(state, step, settings)

    # Update step status to running
    updated_plan = []
    for s in state["plan"]:
        if s["step_id"] == step["step_id"]:
            updated_plan.append({**s, "status": StepStatus.RUNNING,
                                  "started_at": datetime.now(timezone.utc).isoformat()})
        else:
            updated_plan.append(s)

    llm = get_llm("executor", temperature=0.1)

    context = _build_context(state)
    user_msg = (
        f"Execute this step:\n"
        f"Title: {step['title']}\n"
        f"Description: {step['description']}\n"
        f"Preferred tool: {step.get('tool', 'any')}\n\n"
        f"Context:\n{context}"
    )

    last_error = None
    for attempt in range(settings.max_retries):
        try:
            response = llm.invoke(
                [SystemMessage(content=EXECUTOR_SYSTEM), HumanMessage(content=user_msg)],
                tools=TOOL_SCHEMAS,
            )

            tokens = response.usage_metadata.get("total_tokens", 0) if response.usage_metadata else 0

            if not response.tool_calls:
                # LLM gave a text answer (for simple steps)
                result = {"text": response.content, "source": "llm_direct"}
                return _step_success(state, step, updated_plan, result, tokens)

            # Execute the tool call
            tool_call = response.tool_calls[0]
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]

            log.info("Tool call", tool=tool_name, args=str(tool_args)[:100])

            # Run async tool in sync context
            tool_result = asyncio.run(_safe_tool_call(tool_name, tool_args))

            # Log tool call
            tool_log_entry = {
                "step_id": step["step_id"],
                "tool": tool_name,
                "args": tool_args,
                "result_status": tool_result.get("status"),
                "attempt": attempt + 1,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

            if tool_result.get("status") == "error":
                last_error = tool_result.get("error", "Tool error")
                log.warning("Tool failed", tool=tool_name, error=last_error, attempt=attempt+1)
                if attempt < settings.max_retries - 1:
                    # Back off before retrying — important for rate-limited tools like web_search
                    time.sleep(2 ** attempt)  # 1s, 2s, 4s
                    user_msg += f"\n\nAttempt {attempt+1} failed: {last_error}. Try a different approach or query."
                    continue

            return _step_success(
                state, step, updated_plan, tool_result, tokens,
                tool_log=tool_log_entry,
            )

        except Exception as e:
            last_error = str(e)
            log.error("Executor attempt failed", attempt=attempt+1, error=str(e))
            if attempt == settings.max_retries - 1:
                return _step_failed(state, step, updated_plan, last_error)

    return _step_failed(state, step, updated_plan, last_error or "Max retries exceeded")


async def _safe_tool_call(name: str, args: dict) -> dict:
    try:
        return await execute_tool(name, args)
    except Exception as e:
        return {"status": "error", "error": str(e)}


def _handle_synthesize(state: WorkflowState, step: dict, settings) -> WorkflowState:
    """Handle the special synthesize step — produces final answer."""
    llm = get_llm("executor", temperature=0.3)

    context = _build_context(state)
    response = llm.invoke([
        SystemMessage(content=SYNTHESIZE_SYSTEM),
        HumanMessage(content=f"Task: {state['task']}\n\nAll collected results:\n{context}"),
    ])

    final = response.content
    tokens = response.usage_metadata.get("total_tokens", 0) if response.usage_metadata else 0

    updated_plan = [{**s, "status": StepStatus.DONE, "result": "Synthesized"} if s["step_id"] == step["step_id"] else s
                    for s in state["plan"]]

    return {
        **state,
        "plan": updated_plan,
        "final_output": final,
        "step_results": {**state["step_results"], step["step_id"]: {"synthesis": final}},
        "status": TaskStatus.REFLECTING if settings.enable_reflection else TaskStatus.COMPLETED,
        "total_tokens": state["total_tokens"] + tokens,
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "events": state["events"] + [
            make_agent_event(AgentRole.EXECUTOR, "synthesis_complete",
                             f"Final answer synthesized ({len(final)} chars)")
        ],
    }


def _step_success(state, step, updated_plan, result, tokens, tool_log=None):
    final_plan = [{**s, "status": StepStatus.DONE,
                   "result": str(result)[:500],
                   "finished_at": datetime.now(timezone.utc).isoformat()}
                  if s["step_id"] == step["step_id"] else s
                  for s in updated_plan]

    new_tool_log = state["tool_calls_log"] + ([tool_log] if tool_log else [])
    new_events = state["events"] + [
        make_agent_event(AgentRole.EXECUTOR, "step_complete",
                         f"Step '{step['title']}' completed",
                         {"step_id": step["step_id"], "tool": step.get("tool")})
    ]

    # Check if all steps done
    all_done = all(s["status"] in ("done", "skipped") for s in final_plan)
    new_status = (TaskStatus.REFLECTING if get_settings().enable_reflection
                  else TaskStatus.COMPLETED) if all_done else TaskStatus.EXECUTING

    return {
        **state,
        "plan": final_plan,
        "step_results": {**state["step_results"], step["step_id"]: result},
        "tool_calls_log": new_tool_log,
        "status": new_status,
        "total_tokens": state["total_tokens"] + tokens,
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "events": new_events,
    }


def _step_failed(state, step, updated_plan, error):
    final_plan = [{**s, "status": StepStatus.FAILED, "error": error}
                  if s["step_id"] == step["step_id"] else s
                  for s in updated_plan]
    settings = get_settings()
    failed_count = sum(1 for s in final_plan if s["status"] == "failed")

    return {
        **state,
        "plan": final_plan,
        "status": TaskStatus.FAILED if failed_count > 2 else TaskStatus.EXECUTING,
        "error_message": f"Step '{step['title']}' failed: {error}",
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "events": state["events"] + [
            make_agent_event(AgentRole.EXECUTOR, "step_failed",
                             f"Step '{step['title']}' failed after retries: {error}")
        ],
    }