Skip to content

agent graph

Jump to bottom
aakash-anko edited this page Jun 25, 2026 · 3 revisions

agent-graph

Builds and compiles a LangGraph StateGraph with an agent node, a tool node, and conditional routing — the core loop that lets an LLM call codewalk tools in a conversation.

Key Concepts

Term Definition Example
edge A connection between two vertices in a graph, representing a relationship (e.g., an import). If pipeline.py imports scanner.py, there's a directed edge pipeline.py → scanner.py.
embedding A numerical vector (list of numbers) that represents the meaning of text. Similar text → similar vectors. The code def add(a, b): return a+b might become [0.12, -0.45, 0.78, ...] (768 numbers for jinaai/jina-code-embeddings-1.5b via Ollama/MPS).
vector store A database optimized for storing embeddings and finding the most similar ones quickly. ChromaDB stores code chunk embeddings and returns the 5 most similar chunks to your query.
ChromaDB An open-source vector database for storing and searching embeddings. Used here to store code chunks. collection.query(query_texts=["scan files"], n_results=5) returns the 5 closest code chunks.
AST Abstract Syntax Tree — a tree representation of source code structure, where each node is a language construct (function, class, if-statement, etc.). def add(a, b): return a+b becomes a tree: FunctionDef → [args: a, b] → [body: Return → BinOp(a + b)].
LLM Large Language Model — an AI model (like GPT-4, Claude) that generates text given a prompt. get_llm() returns a ChatOpenAI instance that can answer questions about code.

Source: src/codewalk/agent/graph.py

AgentState — line 27

A TypedDict that holds the conversation state flowing through the graph. Has one field: messages, annotated with add_messages so new messages get appended (not replaced).

Example

state = AgentState(messages=[
    SystemMessage(content="You are a codebase expert..."),
    HumanMessage(content="How does auth work?"),
])

state["messages"] → [SystemMessage(...), HumanMessage(...)]

When the agent returns {"messages": [AIMessage(content="Auth uses JWT...")]}, LangGraph's add_messages annotation appends the new AIMessage to the existing list:

state["messages"] → [SystemMessage(...), HumanMessage(...), AIMessage(content="Auth uses JWT...")]

create_agent — line 34

Factory that builds a compiled LangGraph agent with tools, an LLM, conditional routing, and memory (checkpointer). Returns a compiled graph you call with .invoke() or .stream().

Example

Input:

store = VectorStore(persist_dir="/tmp/chroma")
modules_result = {
    "modules": {"api": {"files": ["src/api/routes.py"], "file_count": 1, "languages": {"Python": 1}}},
    "module_graph": {"api": ["config"]},
    "source_root": "src",
    "stats": {"total_files": 10},
}
files = [{"file_path": "src/api/routes.py", "language": "Python"}]
deps = {"graph": {"src/api/routes.py": ["src/config.py"]}}

Step 1 — Create tools (line 50):

tools = create_tools(store, modules_result, files=files, deps=deps, graph_runtime=None, graph_store=None)

→ tools = [<function search_codebase>, <function get_module_info>, ..., <function verify_fix>] (11 tools)

Step 2 — Create LLM with tools bound (line 53–54):

llm = get_llm(temperature=0, reasoning=False)

→ an LLM instance (e.g., ChatOpenAI or ChatAnthropic)

llm_with_tools = llm.bind_tools(tools)

→ same LLM but now aware of all 11 tool schemas — it can emit structured tool_calls

Step 3 — Define agent_node (line 58–87): An inner function that:

  1. Prepends AGENT_SYSTEM_PROMPT to the message history
  2. Calls the LLM
  3. If the LLM outputs a tool call as plain JSON text instead of structured tool_calls, parses and converts it (fallback logic)

Step 4 — Define should_continue (line 90–94): A routing function that checks the last message. If it has tool_calls → return "tools". Otherwise → return END.

Step 5 — Build the graph (line 97–107):

graph = StateGraph(AgentState)
graph.add_node("agent", agent_node)
graph.add_node("tools", ToolNode(tools))
graph.add_edge(START, "agent")
graph.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END})
graph.add_edge("tools", "agent")

This creates the loop: START → agent → (tools → agent)* → END

Step 6 — Compile with memory (line 110–111):

memory = MemorySaver()
return graph.compile(checkpointer=memory)

Returns: a compiled CompiledStateGraph object.

agent_node (inner function) — line 66

Called every time the graph enters the "agent" node. Prepends the system prompt, calls the LLM, and handles a fallback for models that output tool calls as JSON text.

Example

Input state:

state = {"messages": [HumanMessage(content="What does scan_directory do?")]}

Line 68: messages = [SystemMessage(content=AGENT_SYSTEM_PROMPT)] + [HumanMessage(content="What does scan_directory do?")] → [SystemMessage(content="You are a codebase expert..."), HumanMessage(content="What does scan_directory do?")]

Line 69: response = llm_with_tools.invoke(messages)

Case A — LLM returns a structured tool call: → response = AIMessage(content="", tool_calls=[{"name": "explain_function", "args": {"function_name": "scan_directory"}, "id": "call_123"}])

Line 72: response.tool_calls is truthy → skip fallback

Line 86: return {"messages": [response]}

Case B — LLM returns JSON as plain text (fallback): → response = AIMessage(content='I need to look that up. {"name": "explain_function", "arguments": {"function_name": "scan_directory"}}')

Line 72: response.tool_calls is [] (empty) and response.content is truthy → enter fallback

Line 73: match = _TOOL_CALL_RE.search(response.content) → matches the JSON block

Line 75: tool_name = "explain_function"

Line 77: tool_names = {"search_codebase", "get_module_info", "explain_function", "get_overview", "get_blast_radius_map", "get_reading_order", "get_execution_flow", "load_guidelines", "get_architecture_health", "apply_fix", "verify_fix"} (set of 11 names)

Line 78: "explain_function" in tool_names → True

Line 80: tool_args = {"function_name": "scan_directory"}

Line 82–85: response = AIMessage(content="", tool_calls=[{"name": "explain_function", "args": {"function_name": "scan_directory"}, "id": "call_explain_function"}])

Line 86: return {"messages": [response]}

should_continue (inner function) — line 90

Routing function that decides whether the agent should call a tool or end the conversation.

Example

Case A — tool call present:

state = {"messages": [..., AIMessage(content="", tool_calls=[{"name": "explain_function", ...}])]}

Line 92: last_message = state["messages"][-1] → the AIMessage with tool_calls

Line 93: last_message.tool_calls → [{"name": "explain_function", ...}] → truthy

Line 94: return "tools"

The graph routes to the "tools" node (which executes the tool and returns the result).

Case B — final answer (no tool call):

state = {"messages": [..., AIMessage(content="scan_directory walks a directory tree and returns file metadata.")]}

Line 92: last_message.tool_calls → [] → falsy

Line 95: return END

The graph terminates — the AI's text response is the final answer.

_TOOL_CALL_RE — line 61

A compiled regex pattern that matches JSON tool calls embedded in LLM text output. Used by the agent_node fallback logic.

Example

Pattern: \{\s*"name"\s*:\s*"(\w+)"\s*,\s*"arguments"\s*:\s*(\{[^{}]*\})\s*\}

Input text: 'Let me look that up. {"name": "search_codebase", "arguments": {"query": "auth flow"}}'

match = _TOOL_CALL_RE.search(text)
match.group(1)  → "search_codebase"
match.group(2)  → '{"query": "auth flow"}'

Input text: 'The answer is 42.' → match = None (no JSON tool call found, no fallback triggered).

Clone this wiki locally