DEV Community

kai wen ng
kai wen ng

Posted on

Agent Graph from Scratch: Discussion Needed

Building AI agents from scratch allows developers to have full control and customize every workflow.
Recently, in one of my projects, I built an agent with the following workflow:

Workflow Description

The agent begins at the START node, where the planning() function analyzes the user's request and determines the execution path. If the request can be answered immediately, the workflow proceeds to IMMEDIATE_REPLY. If the query requires refinement, it is first processed by CONSTRUCT_NEW_QUERY before continuing. Otherwise, execution moves directly to PLAN_EXECUTION.

During PLAN_EXECUTION, the agent prepares the retrieval strategy, followed by PROJECT_QUERY, which retrieves the relevant project information. The collected information is then used by the END node to generate a response for the user.

After generating the response, the workflow checks whether additional information is required. If so, EMBEDDING_QUERY performs semantic retrieval to gather more context before returning to END to regenerate the response. This retrieval loop continues until no further information is needed or the retry limit is reached, after which the workflow terminates through IMMEDIATE_REPLY.

State Management in Agentic Graphs

One thing I've started noticing while building graph-based agents is how quickly the state grows.
The workflow itself is straightforward: each node performs one task, routing edges decide where to go next, and every node appends its output back into the shared state. However, after a few iterations, the agent accumulates planning decisions, retrieval results, generated queries, intermediate responses, retry counters, and many other pieces of information.
To keep the state manageable, I grouped related fields into smaller Pydantic models instead of placing every attribute directly into one large state.

from pydantic import BaseModel

# Individual state containers
class PlanningState(BaseModel):
    ...
    # Stores routing decisions made during planning

class QueryState(BaseModel):
    ...
    # Stores structured search parameters

class ResponseState(BaseModel):
    ...
    # Stores the generated response and follow-up actions

# Main agent state
class AgentState(BaseModel):
    user_input: str
    chat_history: str | None = None

    planning: PlanningState | None = None
    query: QueryState | None = None
    response: ResponseState | None = None

    retrieved_projects: list | None = None
    retrieved_documents: list | None = None

    retry_count: int = 0
Enter fullscreen mode Exit fullscreen mode

This keeps the top-level state relatively clean while still giving every node a single input/output object.

The concern I have is scalability. As the workflow becomes more sophisticated, the AgentState will inevitably contain more nested states. While accessing data remains simple (state.planning.xxx, state.query.xxx, etc.), it feels like I'm simply moving the complexity one level deeper rather than reducing it.

Another design decision I made is that each node should do exactly one thing. A node performs its task and updates the state—it does not contain if/else branching. All decision making is handled by the graph's routing edges. In other words:

  • Nodes perform work.
  • Edges decide where execution goes next.

This makes each node easier to test and reason about, but it also means the routing logic becomes increasingly complex as the graph grows.

I'm curious how others approach this problem:

  • How do you prevent the shared state from becoming a "God object"?
  • Do you split state by capability, by workflow stage, or something else?
  • Do you keep branching entirely in the graph, or allow nodes to contain some business logic and conditional execution?
  • At what point does a single graph become too large, and when do you split it into subgraphs or sub-agents?

I'd be interested to hear how others balance graph simplicity, state management, and maintainability as agentic workflows become more complex.

Top comments (0)