DEV Community

Cover image for Building an Observable Agent Graph from Scratch
kai wen ng
kai wen ng

Posted on

Building an Observable Agent Graph from Scratch

Observability is important for LLM agents. Understanding the decision process, node execution, latency, and success rate of individual calls becomes increasingly important in a production system.

Instead of simply relying on existing tools such as Langfuse and LangGraph, I wanted to understand what happens underneath and build a customised solution for my own use case.

The Graph

The first version of the graph supports three fundamental execution mechanisms:

  1. Node execution and directed edges
  2. Conditional routing
  3. Fan-out and fan-in execution

This design follows the principle described by Bloch that β€œAll programmers are API designers. Good programs are modular, and intermodular boundaries define APIs. Good modules get reused.” The graph therefore treats each node as an independent computational unit with a well-defined interface to the rest of the graph.

For simplicity, each node is initially restricted to a single outgoing transition. Sequential execution can therefore be represented by returning a single node name, while more complex execution patterns are represented explicitly through routing or fan-out mechanisms.

Each node receives the current graph state and returns a StateUpdate containing the state fields that should be modified. The same principle is applied to parallel execution. A fan-out operation can execute multiple nodes against the same input state, after which their resulting state updates are collected and merged. The fan-in mechanism is responsible for combining these updates while preventing unintended state overwrites.

The graph structure can consequently be represented using a collection of nodes, edges, routers, fan-outs, and fan-ins:

@dataclass
class Edge:
    source: str
    target: str

@dataclass
class FanOut:
    source: str
    targets: list[str]

@dataclass
class FanIn:
    sources: list[str]
    target: str

@dataclass
class StateUpdate:
    values: dict[str, Any] = field(default_factory=dict)
    observation: dict[str, Any] = field(default_factory=dict)

@dataclass
class Node:
    name: str
    func: Callable[[Any], Awaitable[StateUpdate]]
Enter fullscreen mode Exit fullscreen mode

After defining the data classes, we implemented the graph itself with a set of unified functions for adding nodes, edges, routers, fan-out, and fan-in operations. Each function performs type checking and basic validation when the corresponding graph component is added, ensuring that invalid graph structures are detected as early as possible.

class Graph:
    def __init__(self):
        self.nodes: dict[str, Node] = {}
        self.edges: dict[str, Edge] = {}
        self.routers: dict[str, Callable] = {}
        self.fan_out: dict[str, FanOut] = {}
        self.fan_in: dict[str, FanIn] = {}
        self.compiled: CompiledGraph | None = None

    def add_node(self, name, func):
        self.nodes[name] = Node(name, func)

    def add_edge(self, source, target):
        if source in self.edges:
            raise ValueError(f"Node '{source}' already has an edge")
        self.edges[source] = Edge(source, target)

    def add_conditional_edges(self, source, router):
        if source in self.routers:
            raise ValueError(f"Node '{source}' already has a router")
        self.routers[source] = router

    def add_fanout(self, source, targets):
        if source in self.fan_out:
            raise ValueError(f"Node '{source}' already has a fanout")
        self.fan_out[source] = FanOut(source, targets)

    def add_fanin(self, sources, target):
        if target in self.fan_in:
            raise ValueError(f"Target '{target}' already has a fanin")
        self.fan_in[target] = FanIn(sources, target)
Enter fullscreen mode Exit fullscreen mode

Graph Construction

The graph is constructed by registering nodes and their corresponding transitions. Nodes and edges are maintained as dictionaries of their respective dataclass representations. Routers are represented separately as callables because their output depends on the current execution state.
Validation checks the structural correctness of the graph. For example, the graph must contain valid START and END nodes, and every edge must reference nodes that are explicitly defined within the graph. This ensures that both the source and destination nodes of every edge are valid and prevents the execution engine from attempting to traverse undefined nodes at runtime.


    def _validate(self):
        if "START" not in self.nodes:
            raise ValueError("START node is required")

        if "END" not in self.nodes:
            raise ValueError("END node is required")

        for source, edge in self.edges.items():
            if source not in self.nodes:
                raise ValueError(f"Unknown source node: {source}")

            if edge.source != source:
                raise ValueError(
                    f"Edge source mismatch: {source} != {edge.source}"
                )

            if edge.target not in self.nodes:
                raise ValueError(f"Unknown target node: {edge.target}")

        for source, fanout in self.fan_out.items():
            if source not in self.nodes:
                raise ValueError(f"Unknown fan-out source: {source}")

            if fanout.source != source:
                raise ValueError(
                    f"Fan-out source mismatch: {source} != {fanout.source}"
                )

            for target in fanout.targets:
                if target not in self.nodes:
                    raise ValueError(f"Unknown fan-out target: {target}")
Enter fullscreen mode Exit fullscreen mode

Graph Execution

When the compiled graph is executed, the initial state is passed into the graph executor at runtime. The state therefore belongs to an individual graph execution rather than to the LLM agent itself. The graph executor is responsible for determining which graph element should execute next as pass through the _execute_node for execution. To unify the response for each node, every node is expected to return a StateUpdate, except for the terminal END node. An empty values dictionary represents a valid execution in which the node does not modify the graph state.

@dataclass
class CompiledGraph:
    nodes: dict[str, Node]
    edges: dict[str, str]
    routers: dict[str, Callable[[Any], str | list[str]]]
    fanouts: dict[str, list[str]]
    fanins: dict[str, str]
    repo: ObservationRepository
    start: str = "START"
    end: str = "END"
    trace_id: str = None

    async def run(self, state):
        return await self._execute(self.start, state)

    async def _execute_node(self, name, state):
        try:
            node = self.nodes[name]
            start = time.perf_counter()
            update = await node.func(state)
            elapsed = time.perf_counter() - start
            await self._record_observation(
                name=name,
                elapsed=elapsed,
                update=update,
            )
            return update

        except Exception as e:
            await self._record_observation(
                name=name,
                elapsed=time.perf_counter() - start,
                update=StateUpdate(),
                success=False,
                error=str(e),
            )
            raise

    async def _execute(self, name, state):
        if name == self.end:
            return state

        update = await self._execute_node(name, state)
        state = self._apply_update(state, update)

        if name in self.routers:
            targets = self.routers[name](state)

            if isinstance(targets, str):
                return await self._execute(targets, state)

            return await self._execute_fanout(name, targets, state)

        if name in self.fanouts:
            return await self._execute_fanout(name, self.fanouts[name], state)

        edge = self.edges.get(name)

        if edge is None:
            raise ValueError(f"Node '{name}' has no outgoing edge")

        return await self._execute(edge.target, state)

    def _apply_update(self, state, update: StateUpdate):
        for field, value in update.values.items():
            if field not in state.model_fields:
                raise ValueError(f"Unknown state field: '{field}'")
            setattr(state, field, value)

        return state
Enter fullscreen mode Exit fullscreen mode

When a fan-out is encountered, the graph executor resolves the target nodes and executes them through the same execute_node function. The fan-out therefore does not require a second implementation of logging, exception handling, or observation logic.

    async def _execute_fanout(self, source, targets, state):
        results = await asyncio.gather(
            *[self._execute_node(target, state) for target in targets]
        )

        update = self._merge_updates(results)
        state = self._apply_update(state, update)

        fanin = self.fanins.get(source)
        if fanin is None:
            return state

        return await self._execute(fanin.target, state)

    def _merge_updates(self, updates: list[StateUpdate]):
        fields = {}
        observations = []

        for update in updates:
            for field, value in update.values.items():
                if field in fields:
                    raise ValueError(f"State conflict: '{field}'")
                fields[field] = value

            if update.observation is not None:
                observations.append(update.observation)

        return StateUpdate(
            values=fields,
            observation={"fanout": observations} if observations else None,
        )
Enter fullscreen mode Exit fullscreen mode

As all node executions must pass through _execute_node, it becomes straightforward to consistently record what happens during each node execution. This provides a single observation point where we can capture information such as the node name, input state, execution time, returned state update, execution status, and any exceptions that occur. More importantly, this approach ensures that observability is applied uniformly across all nodes without requiring individual nodes to implement their own logging or monitoring logic.

    async def _record_observation(
        self,
        name: str,
        elapsed: float,
        update: StateUpdate | None = None,
        success: bool = True,
        error: str | None = None,
    ):
        def _jsonable(value):
            if hasattr(value, "model_dump"):
                return value.model_dump(mode="json")
            return value
        metadata = update.observation
        metadata = {
            key: _jsonable(value)
            for key, value in metadata.items()
        }

        kwargs = {
            "node": name,
            "duration_ms": elapsed * 1000,
            "success": success,
            "observations": metadata,
            "error": error,
        }

        if self.trace_id is not None:
            kwargs["trace_id"] = self.trace_id

        record = AgentObservation(**kwargs)

        await self.repo.create(record)

        if self.trace_id is None:
            self.trace_id = record.trace_id
Enter fullscreen mode Exit fullscreen mode

This implementation gives me a clearer understanding of how an agent graph can be constructed from relatively simple execution primitives. Although the current implementation is intentionally minimal, the same principles can be extended to more complex agent workflows. Nodes provide the computational units, edges define the execution topology, routers provide conditional transitions, and fan-out/fan-in mechanisms provide parallel execution and state aggregation.

Top comments (0)