backend/agents/planner.py

"""
backend/agents/planner.py

The Planner Agent — the strategic brain of the system.

Responsibilities:
  1. Decompose complex tasks into ordered steps
  2. Assign appropriate tools to each step
  3. Define step dependencies (DAG structure)
  4. Replan if Critic requests changes
  5. Decide when task is complete

Uses Gemini function calling to produce structured JSON plan.
"""
from __future__ import annotations
import json
import uuid
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, AgentRole,
    make_plan_step, make_agent_event,
)
from ..core.logger import get_logger

log = get_logger(__name__)

PLANNER_SYSTEM = """You are a task planner. Decompose the task into 2-4 steps MAX.

Tools: web_search, fetch_url, calculate, run_python, write_file, read_file, get_datetime, synthesize
Always end with a synthesize step.

Return ONLY JSON:
{"thinking":"one line","plan":[{"step_id":"step_1","title":"short","description":"what to do","tool":"tool_name","depends_on":[]}],"estimated_complexity":"low|medium|high"}

Rules: minimal steps, no redundancy, synthesize last, depends_on lists prerequisite step_ids."""


REPLAN_ADDITION = """
Replanning — fix these issues: {critique}
Already done: {completed_steps}
Only plan remaining steps."""


PLANNING_TOOL = {
    "type": "function",
    "function": {
        "name": "submit_plan",
        "description": "Submit the structured task plan",
        "parameters": {
            "type": "object",
            "required": ["thinking", "plan"],
            "properties": {
                "thinking": {"type": "string"},
                "plan": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "required": ["step_id", "title", "description", "tool"],
                        "properties": {
                            "step_id": {"type": "string"},
                            "title": {"type": "string"},
                            "description": {"type": "string"},
                            "tool": {"type": "string"},
                            "depends_on": {"type": "array", "items": {"type": "string"}},
                        },
                    },
                },
                "estimated_complexity": {"type": "string", "enum": ["low", "medium", "high"]},
            },
        },
    },
}


def planner_node(state: WorkflowState) -> WorkflowState:
    """
    LangGraph node — runs the Planner agent.
    Called at start and whenever Critic requests replanning.
    """
    from ..core.config import get_settings
    settings = get_settings()
    llm = get_llm("planner", temperature=0.2)

    log.info("Planner running", task_id=state["task_id"], iteration=state["iteration"])

    # Build prompt — include critique context if replanning
    system_content = PLANNER_SYSTEM
    if state.get("needs_replanning") and state.get("critique"):
        completed = [
            s["title"] for s in state["plan"]
            if s["status"] == "done"
        ]
        system_content += REPLAN_ADDITION.format(
            critique=state["critique"],
            completed_steps="\n".join(f"- {t}" for t in completed) or "None",
        )

    # Memory context
    memory_context = ""
    if state.get("memories"):
        mem_texts = [m["content"] for m in state["memories"][:3]]
        memory_context = "\n\nRelevant past experience:\n" + "\n".join(f"• {t}" for t in mem_texts)

    user_msg = f"Task: {state['task']}{memory_context}"
    if state.get("step_results"):
        results_summary = json.dumps(
            {k: str(v)[:200] for k, v in list(state["step_results"].items())[-3:]},
            indent=2
        )
        user_msg += f"\n\nRecent results:\n{results_summary}"

    try:
        response = llm.invoke(
            [SystemMessage(content=system_content), HumanMessage(content=user_msg)],
            tools=[PLANNING_TOOL],
            tool_choice={"type": "function", "function": {"name": "submit_plan"}},
        )

        tool_call = response.tool_calls[0] if response.tool_calls else None
        if not tool_call:
            raise ValueError("No tool call in planner response")

        plan_data = tool_call["args"]
        thinking = plan_data.get("thinking", "")

        # Convert to plan steps (hard cap)
        new_plan = []
        raw_steps = plan_data.get("plan", [])[:settings.max_plan_steps]
        for step in raw_steps:
            new_plan.append(make_plan_step(
                step_id=step.get("step_id", f"step_{len(new_plan)+1}"),
                title=step.get("title", ""),
                description=step.get("description", ""),
                tool=step.get("tool"),
                depends_on=step.get("depends_on", []),
            ))

        log.info("Plan created", steps=len(new_plan), thinking=thinking[:100])

        # Preserve completed steps if replanning
        if state.get("needs_replanning"):
            completed = [s for s in state["plan"] if s["status"] == "done"]
            final_plan = completed + new_plan
        else:
            final_plan = new_plan

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

        return {
            **state,
            "status": TaskStatus.EXECUTING,
            "plan": final_plan,
            "current_step_index": len([s for s in final_plan if s["status"] == "done"]),
            "needs_replanning": False,
            "iteration": state["iteration"] + 1,
            "total_tokens": state["total_tokens"] + tokens,
            "updated_at": datetime.now(timezone.utc).isoformat(),
            "events": state["events"] + [
                make_agent_event(
                    AgentRole.PLANNER, "plan_created",
                    f"Created {len(new_plan)}-step plan: {thinking[:200]}",
                    {"steps": [s["title"] for s in new_plan], "complexity": plan_data.get("estimated_complexity")},
                )
            ],
        }

    except Exception as e:
        log.error("Planner failed", error=str(e))
        return {
            **state,
            "status": TaskStatus.FAILED,
            "error_message": f"Planner failed: {e}",
            "events": state["events"] + [
                make_agent_event(AgentRole.PLANNER, "error", f"Planning failed: {e}")
            ],
        }