Python Programming Series

Posted on Jan 14

Beyond Linear Chains: Building a Self-Correcting AI Agent with LangGraph.js

#langgraph #langchain #typescript #javascript

Stop building brittle AI pipelines. Learn to build cyclical, stateful agents that check their own work using the new Gemini 2.5 Flash.

If you have built an AI application using standard chains (like A -> B -> C), you have likely encountered the "Happy Path" Fallacy.

It works like this: You ask the LLM to generate data (Step A). You parse that data (Step B). You save it to a database (Step C). It works perfectly 90% of the time. But the moment the LLM hallucinates a format in Step A, Step B crashes, and your entire application throws an error.

In the early days of LLM engineering, the solution was to add more error handling code. Today, the solution is architectural. We are moving from Linear Chains (DAGs) to Cyclical Graphs.

We don't want the process to crash when the AI makes a mistake. We want the AI to notice the mistake, loop back, and fix it.

In this tutorial, we will build a Self-Correcting Agent using LangGraph.js and Google's blazing-fast Gemini 2.5 Flash.

The Architecture: The "Drafter-Critic" Loop

Instead of a straight line, we are building a state machine with a cycle.

The Drafter (Node A): Generates an initial answer.
The Critic (Node B): Evaluates the answer. Is it accurate? Does it follow constraints?
The Router (Edge):
- If the quality is Good → End.
- If the quality is Bad → Loop back to Node A (with feedback).

This loop mimics how human developers work. We don't write perfect code on the first try; we write, review, debug, and rewrite.

Prerequisites

You will need a Node.js project setup.

Get a Gemini API Key: Head to Google AI Studio and grab a key.
Install Dependencies:

npm install @langchain/google-genai @langchain/langgraph @langchain/core dotenv zod

Step 1: Define the Graph State

In LangGraph, the State is the shared memory of your application. Every node (agent) reads from this state and writes updates to it.

import { Annotation, MessagesAnnotation } from "@langchain/langgraph";
import { BaseMessage } from "@langchain/core/messages";

// We extend the default MessagesAnnotation to include a 'feedback' field
// and a 'loopCount' to prevent infinite loops.
const AgentState = Annotation.Root({
  ...MessagesAnnotation.spec,
  feedback: Annotation<string>,
  loopCount: Annotation<number>,
});

Step 2: Initialize Gemini 2.5 Flash

We use the Google GenAI adapter. Note that we are specifying the latest gemini-2.5-flash model, which offers improved reasoning capabilities at high speed—perfect for iterative loops where latency matters.

import { ChatGoogleGenerativeAI } from "@langchain/google-genai";
import * as dotenv from "dotenv";
dotenv.config();

const model = new ChatGoogleGenerativeAI({
  modelName: "gemini-2.5-flash",
  apiKey: process.env.GOOGLE_API_KEY,
  temperature: 0.7,
});

Step 3: Create the Nodes

We need two distinct functions (Nodes): one to generate content, and one to critique it.

The Drafter Node

This node looks at the state. If there is feedback from a previous attempt, it incorporates it.

async function drafterNode(state: typeof AgentState.State) {
  const { messages, feedback } = state;
  const lastMessage = messages[messages.length - 1];
  const loopCount = state.loopCount || 0;

  console.log(`\n--- DRAFTER (Attempt ${loopCount + 1}) ---`);

  // If this is a retry, add the critique to the prompt
  let prompt = "";
  if (feedback) {
    prompt = `Your previous attempt received this feedback: "${feedback}". 
    Please rewrite the answer to address this. 
    Original Request: ${messages[0].content}`;
  } else {
    // First attempt
    prompt = `You are a helpful assistant. Please answer the user's request: ${messages[0].content}`;
  }

  const response = await model.invoke(prompt);

  // Return the update to the state
  return { 
    messages: [response], 
    loopCount: loopCount + 1 
  };
}

The Critic Node

This node analyzes the Drafter's output. It acts as a quality gate.

import { z } from "zod";

async function criticNode(state: typeof AgentState.State) {
  const { messages } = state;
  const lastResponse = messages[messages.length - 1].content;

  console.log("\n--- CRITIC IS REVIEWING ---");

  // We ask Gemini to act as a judge
  const prompt = `
    Review the following text for brevity and clarity. 
    If it is over 30 words, it is REJECTED.

    TEXT: "${lastResponse}"

    Respond ONLY with "ACCEPT" or "REJECT: [Reason]".
  `;

  const response = await model.invoke(prompt);
  const decision = response.content as string;

  if (decision.includes("REJECT")) {
    return { feedback: decision };
  }

  return { feedback: "APPROVED" };
}

Step 4: The Conditional Edge (The Logic)

This function determines the flow of the graph. It doesn't modify state; it just points to the next node.

import { END } from "@langchain/langgraph";

function shouldContinue(state: typeof AgentState.State) {
  const { feedback, loopCount } = state;

  // Safety break: Stop after 3 attempts to save tokens
  if (loopCount && loopCount >= 3) {
    console.log("--- MAX ITERATIONS REACHED ---");
    return END;
  }

  if (feedback === "APPROVED") {
    console.log("--- QUALITY CHECK PASSED ---");
    return END;
  }

  console.log("--- QUALITY CHECK FAILED: RETRYING ---");
  return "drafter"; // Loop back!
}

Step 5: Wiring the Graph

Finally, we assemble the pieces.

import { StateGraph } from "@langchain/langgraph";

const workflow = new StateGraph(AgentState)
  .addNode("drafter", drafterNode)
  .addNode("critic", criticNode)
  .addEdge("__start__", "drafter") // Start here
  .addEdge("drafter", "critic")    // Drafter always goes to Critic
  .addConditionalEdges("critic", shouldContinue, {
    drafter: "drafter", // The loop
    [END]: END          // The exit
  });

const app = workflow.compile();

Running the Agent

Let's test it with a prompt that forces the AI to be verbose, triggering the Critic's length constraint loop.

async function main() {
  const inputs = {
    messages: [
      {
        role: "user",
        content: "Explain Quantum Entanglement in extreme detail.",
      },
    ],
    loopCount: 0,
  };

  const result = await app.invoke(inputs);

  console.log("\nFINAL OUTPUT:");
  console.log(result.messages[result.messages.length - 1].content);
}

main();

What happens when you run this?

Pass 1: Gemini 2.5 Flash explains Quantum Entanglement in 200 words.
Critic: Sees 200 words. Checks rule (>30 words). Rejects it.
Router: Sees "REJECT". Routes back to "Drafter".
Pass 2: Drafter sees the feedback: "Too long". It rewrites it to be shorter.
Critic: Sees 25 words. Approves.
End: You get a concise, validated answer.

Why this changes everything

This simple pattern—Generation, Reflection, and Iteration—is the secret sauce behind today's most advanced AI tools.

By using LangGraph.js for orchestration and Gemini 2.5 Flash for low-latency inference, you have built a system that is resilient. It doesn't just crash when the output is wrong; it fixes itself.

This is no longer just a script; it is an Autonomous Agent.

Want to learn more?

This article covers the basics of cyclical graphs. To master the full architecture of production-grade AI applications, check out my book:

📘 AI with JavaScript & TypeScript: Foundations

Relevant Chapters for this topic:

Chapter 16: Building Agents with LangGraph.js (Deep dive into state machines).
Chapter 17: Multi-Agent Systems in Node.js (How to make multiple agents collaborate).
Chapter 8: Function Calling with TypeScript (Giving your agents tools to interact with the real world).

Stop guessing. Start engineering.

Explore also the complete multi-volume "Python Programming Series" for a comprehensive journey from Python fundamentals to advanced AI deployment: https://www.amazon.com/dp/B0FTTQNXKG . Each book can be read as a standalone.