richiejp logo

Using small LLMs in agents

What usually gets the headline attention with agents is when they code a website or create a research report autonomously. This is usually done with gigantic frontier models that only a few companies can run.

These models usually have 400 billion parameters or more, possibly even a trillion. In the case of the proprietary models it’s kept secret how many parameters they have and even if you could afford the hardware to run them you are not allowed to.

Open Source models you can practically run on commercial hardware go up to about 70 billion parameters (ignoring MoE). These still have impressive capabilities, sometimes beating the proprietary models in benchmarks.

However the hardware needed to run them is somewhat expensive and the speed at which the model outputs text is relatively slow. The performance is fine for chatting, but once you add in tool calls and reasoning, it can start to seem sluggish.

Meanwhile towards the bottom end of the scale we have models in the 1 to 4 billion parameter range. OpenAI’s GPT-2 was in this range at 1.5 billion parameters.

Models of this size can be executed quickly even on cheap hardware and edge devices. The only catch being many consider them to be toys and rather useless.

However models in this size can still create coherent sentences in response to some input. What’s more they can understand the intent of a command written in plain language and translate that into code.

A traditional use of machine learning is to classify some text and perform a given action depending on what category the text falls into.

For example traders would train classifiers to decide if some news was positive or negative for a stock. Typically the classifier would output a good, bad or neutral rating for a chunk of text fed into it.

As an alternative to training classifiers an LLM can be used. Typically this won’t be as fast, but it is far more flexible and if we are using a small LLM then the performance is still good.

To create the classifier we just need to describe to the LLM what classes it should sort the news into. We don’t have to perform a training step and we can tweak it at any time just by modifying the prompt.

Getting the LLM to respond with just the class name can be tricky, but as I explain further down, there is way to force it to produce valid output using tool calling.

Another example, and what prompted this article, is quickly deciding if a chat bot or agent should respond to a message. The LLM used to power an agent’s reasoning could be very large and expensive. If an agent is in a high volume chat group, we don’t want it to process messages that it can’t help with.

I recently added a feature to LocalAGI called Filters & Triggers to handle this scenario. Presently there are two types of filter, one is good ol’ fashioned Regex and the other is a LLM based classifier.

To implement the classifier I used a trick where we force the LLM to fill out the parameters to a particular tool call.

Below is the Go code which implements the classifier and uses the tool call trick. The GenerateTypedJSON function makes a call to the OpenAI compatible completions API with the jsonschema.Definition to describe the tool call parameters. It specifies one tool is available and that the tool must be used.

const fmtT = `
  Does the below message fit the description "%s"

  %s
  `

func (f *ClassifierFilter) Apply(job *types.Job) (bool, error) {
  input := extractInputFromJob(job)
  guidance := fmt.Sprintf(fmtT, f.description, input)
  var result struct {
    Asserted bool `json:"answer"`
  }
  err := llm.GenerateTypedJSON(job.GetContext(), f.client, guidance, f.model, jsonschema.Definition{
    Type: jsonschema.Object,
    Properties: map[string]jsonschema.Definition{
      "answer": {
        Type:        jsonschema.Boolean,
        Description: "The answer to the first question",
      },
    },
    Required: []string{"answer"},
  }, &result)
  if err != nil {
    return false, err
  }

  if result.Asserted {
    return f.allowOnMatch, nil
  }
  return !f.allowOnMatch, nil
}

We never actually make the tool call, we are just forcing the LLM to fill out the parameters.

When LocalAI receives this request it will check that each token the LLM outputs fits the above schema. So we always get a valid JSON object with a boolean field called answer.

The classification may still be wrong, but we at least get a well formed response. We can generalise this to add more classes to our classifier or, in theory, to force the LLM to produce code other than JSON.

A problem with all LLMs, but especially the smaller variety, is that they can start getting confused after a few steps into a multi-stage plan and start going wildly astray. If they can be forced to follow a flow chart that puts them on rails, then while not every decision may be correct, they at least stick to a well defined path.

So for example, you can have an LLM agent running on a sound bar with voice activation. When the user asks to increase the volume, the LLM can first be forced to classify if it is being spoken too. Then if it chooses yes, you force it to classify what command is being asked of it and if it decides that it is volume, then it can be forced to decide if it is clear how much to change the volume by and if clarification is required.

A device like a sound bar won’t have a large GPU in it, but a small TPU (Tensor Processing Unit) is possible and could run a small LLM. More likely though the bar can be connected to a home hub that has the necessary compute. In any case a small LLM is all that is required for this so long as we keep it on rails. Meanwhile if you want to do this without an LLM, then you will have to create several classifiers and only use canned responses.