Planning

import operator
from typing import Annotated, Literal, TypedDict

from langchain_ollama import ChatOllama
from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel

from docsassistant.search import search_docs

planner_llm = ChatOllama(model="qwen3:8b", reasoning=True, temperature=0)   # a thinking model plans
worker_llm = ChatOllama(model="llama3.1:8b", temperature=0)                 # a small model executes
MAX_REPLANS = 2


class Plan(BaseModel):
    steps: list[str]


class Replan(BaseModel):
    action: Literal["continue", "replan", "respond"]
    new_steps: list[str] = []          # only for "replan"
    reason: str


class ReplanState(TypedDict):
    question: str
    plan: list[str]
    done: Annotated[list[str], operator.add]    # "step -> outcome" records
    replans: int
    answer: str


def plan(state: ReplanState) -> dict:
    p = planner_llm.with_structured_output(Plan).invoke(
        f"Plan documentation lookups (max 4) to answer: {state['question']}")
    return {"plan": p.steps or [state["question"]], "replans": 0}   # never an empty plan


def execute(state: ReplanState) -> dict:
    step, *rest = state["plan"]
    chunks = search_docs(step, k=3)
    outcome = "NO RESULTS" if not chunks else worker_llm.invoke(f"{step}\n{chunks}").content
    return {"plan": rest, "done": [f"{step} -> {outcome[:300]}"]}


def review(state: ReplanState) -> dict:
    decision = planner_llm.with_structured_output(Replan).invoke(
        f"Question: {state['question']}\nDone so far:\n" + "\n".join(state["done"]) +
        f"\nRemaining plan: {state['plan']}\n"
        "Continue, replace the remaining plan (replan), or respond now?")
    if decision.action == "replan" and decision.new_steps and state["replans"] < MAX_REPLANS:
        return {"plan": decision.new_steps, "replans": state["replans"] + 1}
    if decision.action == "continue" and state["plan"]:
        return {}
    answer = worker_llm.invoke("Answer from these findings:\n" + "\n".join(state["done"]))
    return {"plan": [], "answer": answer.content}


def next_node(state: ReplanState) -> str:
    return END if state.get("answer") else "execute"


builder = StateGraph(ReplanState)
builder.add_node("plan", plan)
builder.add_node("execute", execute)
builder.add_node("review", review)
builder.add_edge(START, "plan")
builder.add_edge("plan", "execute")
builder.add_edge("execute", "review")
builder.add_conditional_edges("review", next_node, ["execute", END])
replanner = builder.compile()

import java.util.List;

import dev.langchain4j.agentic.AgenticServices;
import dev.langchain4j.agentic.UntypedAgent;
import dev.langchain4j.agentic.planner.Action;
import dev.langchain4j.agentic.planner.AgentInstance;
import dev.langchain4j.agentic.planner.InitPlanningContext;
import dev.langchain4j.agentic.planner.Planner;
import dev.langchain4j.agentic.planner.PlanningContext;

public class ReviewingPlanner implements Planner {

    private static final int MAX_ROUNDS = 3;                  // the re-plan budget
    private AgentInstance searcher, reviewer, answerer;
    private int rounds = 0;

    @Override
    public void init(InitPlanningContext context) {
        List<AgentInstance> agents = context.subagents();     // in subAgents(...) order
        searcher = agents.get(0);
        reviewer = agents.get(1);
        answerer = agents.get(2);
    }

    @Override
    public Action firstAction(PlanningContext context) {
        return call(searcher);
    }

    @Override
    public Action nextAction(PlanningContext context) {
        String last = context.previousAgentInvocation().agentName();
        if (last.equals(searcher.name())) {
            return call(reviewer);                            // review after every step
        }
        if (last.equals(reviewer.name())) {
            String decision = context.agenticScope().readState("decision", "respond");
            boolean more = decision.equals("continue") && ++rounds < MAX_ROUNDS;
            return more ? call(searcher) : call(answerer);
        }
        return done();                                        // the answerer has written "answer"
    }
}

UntypedAgent replanner = AgenticServices.plannerBuilder()
        .subAgents(docsSearcher, stepReviewer, docsAnswerer)  // @Agent interfaces from agentBuilder(...)
        .outputKey("answer")
        .planner(ReviewingPlanner::new)                       // a fresh planner (and budget) per run
        .build();

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;

public class Replanner {

    record Plan(List<String> steps) {}
    record Review(String action, List<String> newSteps, String reason) {}   // continue | replan | respond

    private static final int MAX_REPLANS = 2;
    private final ChatClient planner;                         // qwen3:8b with thinking
    private final ChatClient worker;                          // llama3.1:8b
    private final VectorStore docs;

    public Replanner(ChatClient planner, ChatClient worker, VectorStore docs) {
        this.planner = planner;
        this.worker = worker;
        this.docs = docs;
    }

    public String answer(String question) {
        Deque<String> plan = new ArrayDeque<>(planner.prompt()
                .user("Plan documentation lookups (max 4) to answer: " + question)
                .call().entity(Plan.class).steps());
        List<String> done = new ArrayList<>();
        int replans = 0;
        while (!plan.isEmpty()) {                             // an empty plan goes straight to the answer
            String step = plan.removeFirst();
            List<Document> chunks = docs.similaritySearch(SearchRequest.builder().query(step).topK(3).build());
            String outcome = chunks.isEmpty() ? "NO RESULTS"
                    : worker.prompt().user(step + "\n" + chunks).call().content();
            done.add(step + " -> " + outcome);
            Review review = planner.prompt()
                    .user("Question: " + question + "\nDone so far:\n" + String.join("\n", done)
                            + "\nRemaining plan: " + plan
                            + "\nContinue, replace the remaining plan (replan), or respond now?")
                    .call().entity(Review.class);
            if (review.action().equals("replan") && review.newSteps() != null && replans < MAX_REPLANS) {
                plan = new ArrayDeque<>(review.newSteps());
                replans++;
            } else if (!review.action().equals("continue")) {
                break;                                        // respond, or the re-plan budget is spent
            }
        }
        return worker.prompt().user("Answer from these findings:\n" + String.join("\n", done)).call().content();
    }
}

from enum import Enum
from typing import Callable


class Status(Enum):
    SUCCESS = "success"
    FAILURE = "failure"
    RUNNING = "running"


class Node:
    def tick(self, bb: dict) -> Status: ...


class Condition(Node):
    def __init__(self, check: Callable[[dict], bool]):
        self.check = check

    def tick(self, bb):
        return Status.SUCCESS if self.check(bb) else Status.FAILURE


class Action(Node):
    def __init__(self, run: Callable[[dict], Status]):
        self.run = run

    def tick(self, bb):
        return self.run(bb)


class Sequence(Node):                       # AND: stop at the first child that does not succeed
    def __init__(self, *children: Node):
        self.children = children

    def tick(self, bb):
        for child in self.children:
            status = child.tick(bb)
            if status is not Status.SUCCESS:
                return status
        return Status.SUCCESS


class Selector(Node):                       # OR: stop at the first child that does not fail
    def __init__(self, *children: Node):
        self.children = children

    def tick(self, bb):
        for child in self.children:
            status = child.tick(bb)
            if status is not Status.FAILURE:
                return status
        return Status.FAILURE