I recently had a realization about the structure of typical llm computations like the agent loop or the llm inference itself which I found interesting to better understand what's going on.

The basic version of an agent loop (the loop that produces an llm conversation) looks something like this, including tool calls:

# start with the system prompt in the conversation history
h = [p]
while True:
   # get user input
   u = userInput()

   # add to context, compute response
   h <- h :: u
   l = llm(h) # call to an OpenAI endpoint for example

   # deal with tool calls as long as there are some
   while l has tool call:
     r = execute tool call
     h <- h :: r
     l = llm(h)

It's pretty straightforward, but it's also a bit inelegant, because you update the history and call the llm in two different places. Sure, you could try to clean it up a bit, and rearrange, but this is closest to "you call an llm, and execute tool calls as long as it keeps coming back with more tool requests".

And it's a bit tricky. I had one subtle bug once where I forgot the second history update, and then wondered why the llm was calling the tool call over and over again! I was already blaming the stupid LLMs, when I didn't realize that to them it just looked like the result never came back! Totally "it's definitely a compiler bug territory" 😅

OK, but what's the structure of the loop?

• we start with a history that's just the system prompt [p]

• depending on the type of the last message we do one of the following...

  • if it's a llm reply and it's a tool call, execute a tool call, and we get a tool result
  • if it's a user message or a tool result, call the llm, get a reply
  • else ask the user for the next input

• ... then we append the result to the history and start over.

This reminded me a lot of functional operations on data streams in languages like Scala, clojure, or Haskell!

A simpler example is the fold or reduce, where you apply a function f to a collection like [a b c] by first applying f(a, b) to get x, and then f(x, b) to get y, the result:

(fold f [a b c]) = (f(f a b) c)

What we have here is a close cousin, the unfold. While a fold summarizes a collection, the unfold generates a collection. The unfold takes a function f and a starting state. f takes a state and either returns an output and a new state, or None when it is done.

For our agent loop, this function looks something like this:

def next(h):
   if last(h) is user message or tool result:
      r = llm(h)
   elif last(h) is tool call:
      r = execute tool call(last(h))
   else:
      r = userInput()
   return (r, h :: r)

The whole agent loop becomes essentially unfold(next, [p]).

So this is at the core of any agent framework out there, no matter how complicated. They add things that aren't here like persistence, streaming responses, retries, interrupts, limits on number of turns, etc., but at the core, it's still an unfold of the next over an append only history!

Another interesting parallel is that the way how the llm interacts with the worlds is reminiscent of how that is modeled in Haskell with IO monads. In Haskell as a pure functional programming language, the language itself has no concept of IO, and IO operations are just descriptions of IO which are then executed by the Haskell runtime. Here, the situation is the same, as the llm cannot make a system call, but it can ask the runtime, the agent loop to do a tool call on its behalf.

That clarifies how each of the components behave when you think about testing and replays:

• The llm itself is almost functionally pure, except for the sampling. If you take a more "probabilistic view" then the llm part can be replayed and will sample from the same distribution, which could still have forks, however.
• You need to "undo" the result of tool calls and create exactly the same situation as before, then you can compare outcomes.
• The human itself is also not "functionally" pure. But if you use an LLM generated conversation partner you can reset that in the same way.

Ok, let's... unfold... the next layer, llm generation

Turns out, the inference part of the llm computation is also this kind of unfolding a computation over an append only history. Very simplified, in particular leaving out the details of RoPE positional embeddings, llm inference looks like this:

llm(h) = unfold(inference, prefill(h))

# prefill computes k, v vectors for all heads / all layers for the context
# which are needed for the attention mechanisms later
def prefill(h):
  embed tokens for the whole history h
  # compute transformer for the whole history from bottom to top
  from bottom layer to top:
    # for each layer we can compute in parallel
    for all tokens in h:
       compute k, q, v vectors
       compute attention to previous positions
       compute mlp
  return [h,
          [k, v] vectors from all layers for all of the history]

# Do one step of inference, update state by appending the
# outputs and [k, v] vectors
def inference(c: [outputs, [k,v] vectors):
  embed last token
  from bottom layer to top:
    compute k, q, v vectors
    compute attention to previous positions
    compute mlp
  sample output o from last layer
  if o is <END>:
    return None
  else:
    return (o, c :: [o, [k, v] vectors ])

So in essence, prefill is like a layer by layer map operation that works in parallel over the whole prefix (and is cached so you only need to compute for the new part - that's the reason why prefill is faster than generation), and inference is an unfold that generates layer by layer.

What is the conversation history in the agent loop are the [k, v] vectors that are required for the attention computation, but structurally it's again an unfold!

If you look at agent libraries that are written in functional languages, you will actually see these patterns in production, for example https://docs.outskirtslabs.com/ol.llx/next/. I actually don't think we should now rewrite all libraries in functional style. Sometimes a for loop is clearer in its intent. But then you add a ton of other stuff and it becomes muddy. But essentially, it's an unfold over an append only history!