LangGraph StateGraph Node & Edge Configuration Errors

from typing import Annotated, TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage


class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    question: str


def retrieve(state: State) -> dict:
    # ... call vector store here ...
    docs = ["doc1", "doc2"]
    return {"messages": [("system", f"context: {docs}")]}


def grade_documents(state: State) -> dict:
    # ... judge relevance ...
    return {"messages": [("system", "grade: relevant")]}


def generate(state: State) -> dict:
    # ... call the LLM ...
    return {"messages": [("assistant", "final answer")]}


# 1. Register EVERY node first
graph = StateGraph(State)
graph.add_node("retrieve", retrieve)
graph.add_node("grade_documents", grade_documents)   # <-- add BEFORE the edge below
graph.add_node("generate", generate)

# 2. Then wire the edges
graph.add_edge(START, "retrieve")
graph.add_edge("retrieve", "grade_documents")
graph.add_edge("grade_documents", "generate")
graph.add_edge("generate", END)

# 3. Compile — this is where the ValueError would fire if a node is missing
app = graph.compile()

from langchain_core.messages import AIMessage

# ❌ Bug: returns None on the "no docs" branch
def bad_node(state):
    if state["docs"]:
        response = llm.invoke(state["messages"])
        return {"messages": [response]}
    # implicit return None → InvalidUpdateError at runtime

# ❌ Bug: returns the AIMessage directly, not a dict
def bad_node_2(state):
    response = llm.invoke(state["messages"])
    return response   # → InvalidUpdateError: Expected dict

# ✅ Correct: every branch returns a dict
def good_node(state):
    if not state["docs"]:
        return {}                              # no state change, but valid
    response = llm.invoke(state["messages"])
    return {"messages": [response]}            # always dict

# ✅ Multi-key update
def planning_node(state):
    plan = build_plan(state["question"])
    return {"plan": plan, "step_index": 0}     # dict of state keys → new values

# ✅ Using Command for state update + explicit routing (LangGraph 0.2+)
from langgraph.types import Command
def router_node(state) -> Command:
    if state["is_math"]:
        return Command(update={"tool": "calculator"}, goto="tool_use")
    return Command(update={"tool": None}, goto="generate")

from enum import Enum
from langgraph.graph import StateGraph, START, END

class N(str, Enum):
    """Frozen enum of node names — typos become NameError at import."""
    RETRIEVE = "retrieve"
    GRADE = "grade_documents"
    GENERATE = "generate"
    REWRITE = "rewrite_query"


def build_graph(State):
    g = StateGraph(State)
    g.add_node(N.RETRIEVE, retrieve)
    g.add_node(N.GRADE, grade_documents)
    g.add_node(N.GENERATE, generate)
    g.add_node(N.REWRITE, rewrite_query)

    g.add_edge(START, N.RETRIEVE)
    g.add_edge(N.RETRIEVE, N.GRADE)
    g.add_conditional_edges(
        N.GRADE,
        lambda s: N.GENERATE if s["relevant"] else N.REWRITE,
        {N.GENERATE: N.GENERATE, N.REWRITE: N.REWRITE},
    )
    g.add_edge(N.REWRITE, N.RETRIEVE)
    g.add_edge(N.GENERATE, END)

    # Optional: assert graph is well-formed before compile
    declared = {N.RETRIEVE.value, N.GRADE.value, N.GENERATE.value, N.REWRITE.value}
    referenced = {e for pair in g.edges for e in pair if e not in (START, END)}
    missing = referenced - declared
    assert not missing, f"Edges reference undeclared nodes: {missing}"

    return g.compile()