backend/agents/orchestrator.py

"""
backend/agents/orchestrator.py

LangGraph workflow graph — connects all agents with conditional routing.

Graph structure:
  START
    │
    ▼
  memory_retrieve         ← Pull relevant past memories
    │
    ▼
  planner                 ← Decompose task into steps
    │
    ▼
  executor ─────────────┐ ← Execute one step at a time
    │                   │   (loops until all steps done)
    │ all steps done    │
    ▼                   │
  critic ────────────────┘ ← Evaluate quality
    │ approve               │ needs_replanning → back to planner
    ▼
  memory_store            ← Save learnings
    │
    ▼
  END

Conditional edges:
  executor → executor      (more steps to run)
  executor → critic        (all steps done)
  critic   → planner       (needs replanning)
  critic   → memory_store  (approved)
"""
from __future__ import annotations
from typing import Literal

from langgraph.graph import StateGraph, START, END

from ..state.graph_state import WorkflowState, TaskStatus
from .planner import planner_node
from .executor import executor_node
from .critic import critic_node
from .memory_agent import memory_retrieve_node, memory_store_node
from ..core.logger import get_logger

log = get_logger(__name__)


# ── Routing functions ─────────────────────────────────────────────────────────

def route_after_executor(state: WorkflowState) -> Literal["executor", "critic", "end"]:
    """
    After executor runs one step, decide what's next:
      - More pending steps → run executor again
      - All done → send to critic
      - Fatal failure → end
    """
    if state["status"] == TaskStatus.FAILED:
        log.warning("Executor fatal failure, ending", task_id=state["task_id"])
        return "end"

    if state["status"] == TaskStatus.REFLECTING:
        return "critic"

    if state["status"] == TaskStatus.COMPLETED:
        return "critic"

    # Check for any pending steps
    pending = [s for s in state["plan"] if s["status"] == "pending"]
    if pending:
        return "executor"

    return "critic"


def route_after_critic(state: WorkflowState) -> Literal["planner", "memory_store"]:
    """
    After critic evaluates:
      - needs_replanning → back to planner (with critique context)
      - approved → store memory and complete
    """
    if state.get("needs_replanning") and state["status"] == TaskStatus.PLANNING:
        log.info("Critic requested replan", task_id=state["task_id"])
        return "planner"
    return "memory_store"


def route_after_planner(state: WorkflowState) -> Literal["executor", "end"]:
    """After planner creates plan, check it's valid before executing."""
    if state["status"] == TaskStatus.FAILED:
        return "end"
    if not state.get("plan"):
        log.error("Planner returned empty plan")
        return "end"
    return "executor"


# ── Graph builder ─────────────────────────────────────────────────────────────

def build_workflow() -> StateGraph:
    """Build and compile the LangGraph workflow."""
    graph = StateGraph(WorkflowState)

    # ── Add nodes ──
    graph.add_node("memory_retrieve", memory_retrieve_node)
    graph.add_node("planner", planner_node)
    graph.add_node("executor", executor_node)
    graph.add_node("critic", critic_node)
    graph.add_node("memory_store", memory_store_node)

    # ── Add edges ──

    # Entry: always start with memory retrieval
    graph.add_edge(START, "memory_retrieve")

    # Memory → Planner
    graph.add_edge("memory_retrieve", "planner")

    # Planner → Executor or END (if plan failed)
    graph.add_conditional_edges(
        "planner",
        route_after_planner,
        {"executor": "executor", "end": END},
    )

    # Executor → Executor (more steps) or Critic (done) or END (fatal)
    graph.add_conditional_edges(
        "executor",
        route_after_executor,
        {"executor": "executor", "critic": "critic", "end": END},
    )

    # Critic → Planner (replan) or Memory store (done)
    graph.add_conditional_edges(
        "critic",
        route_after_critic,
        {"planner": "planner", "memory_store": "memory_store"},
    )

    # Memory store → END
    graph.add_edge("memory_store", END)

    return graph.compile()


# ── Singleton compiled graph ──────────────────────────────────────────────────

_workflow = None


def get_workflow():
    global _workflow
    if _workflow is None:
        _workflow = build_workflow()
        log.info("Workflow graph compiled")
    return _workflow


# ── Run function ──────────────────────────────────────────────────────────────

async def run_workflow(initial_state: WorkflowState) -> WorkflowState:
    """
    Execute the full multi-agent workflow.
    Returns the final state with all results, events, and metrics.
    """
    workflow = get_workflow()

    log.info("Workflow starting",
             task_id=initial_state["task_id"],
             task=initial_state["task"][:80])

    try:
        # LangGraph invoke — runs the full graph
        final_state = await workflow.ainvoke(initial_state)
        log.info("Workflow completed",
                 task_id=initial_state["task_id"],
                 status=final_state.get("status"),
                 score=final_state.get("quality_score"),
                 tokens=final_state.get("total_tokens"))
        return final_state

    except Exception as e:
        log.error("Workflow crashed", error=str(e), task_id=initial_state["task_id"])
        return {
            **initial_state,
            "status": TaskStatus.FAILED,
            "error_message": f"Workflow crashed: {e}",
        }


async def stream_workflow(initial_state: WorkflowState):
    """
    Stream workflow events as they happen.
    Yields state snapshots after each node execution.
    Used by SSE endpoint for real-time UI updates.
    """
    workflow = get_workflow()

    async for chunk in workflow.astream(initial_state, stream_mode="values"):
        yield chunk