Building an AI Agent with a Small Language Model

An AI agent is a system that can carry out a task with a degree of autonomy. It typically combines a large language model, such as Claude or ChatGPT, with external tools that allow it to interact with its wider environment. A software wrapper, often called a harness, can enforce task-specific logic and coordinate how the model uses those tools. Coding assistants are a common example: their harness may run several models at once, with each model focusing on a different part of the task.

This article describes an experiment in building a self-hosted AI agent with a small language model. It explains the test case, the main reliability problems encountered, and the techniques that improved tool use.

Model Size

A small language model AI agent is a loose and time dependant definition, for this article I mean a model with 3 or 4 billion parameters. As of writing, cutting edge models have greater than 100 billion parameters. These large models are very capable and can do a lot. The problem is that, roughly speaking, each billion parameters of a model take up about 1GB of memory.  Computers that can comfortably run 100B+ parameter models are very expensive.  Small models will run on most relatively new laptops. The trade off it that capability scales with the number of parameters. Small models are more challenging to work with.

Test Case

My test case is a chat bot agent that does an AI automation opportunity analysis with clients. This is a repetitive task which is often done in person via a meeting with each client. We can both save time having it done before hand. The threshold for success is if it can be done

  1. correctly,
  2. with a better user experience than filling out a large form.

I’m using llama3.2 from ollama. My AI agent harness is in python.

Tool Use

Using python makes the task of using tools easy as you can provide ollama chat interface a list of tools it has available with prompt like comments describing what the function does. The issue with the small model was that the tool was never called at the right time.

The way around this is to explicitly check if the tool is needed. Thus, my agent had two interactions with my model.

  1. The main conversation with the user based on my initial prompt.
  2. A second conversation with the LLM asking the primary conversation contains certain information.

JSON formatting

For my python harness to interact with the model it needs the output to be in a consistent format. I chose JSON as it’s a widely used format which the language model will have seen plenty of examples of. Asking the model to return its answers in JSON format via the prompt only had limited success. The first response was normally valid, but the second question had extra text which caused issues. Luckily, this is a common enough use case that there is a flag that can be passed to model which forces a JSON return value.

Varied Responses

LLM by their nature are probabilistic. It’s a feature, rather than a bug, that they give different responses to the same question – its how they seem like they are having a genuine conversation and can produce creative output.

However, for specific questions it can cause problems. The level of determinism is controlled via the temperature variable. The next word in each sentence is picked by selecting from the most likely options. The higher the temperature variable the more likely the lower probability word is selected. By setting a low temperature, it reduces the variability in outcomes a lot.

response = chat(
     model=MODEL,
     messages=[system_message],
     options={'temperature': 0.0}
)

Single Word Names

One of the piece of information I wanted to get was the clients name. The interaction normally took the form of a direct question and then response from the user. One issue I found quickly was single word names, e.g. Dave rather than Dave Smith. In my tool query the name wasn’t recognized as one (although it was recognized in the main conversation). This part of the problem was improved by giving the prompt examples, like the ones below:

Input: John -> Output: {'success' : True, 'info_value' : 'John'}
Input: Table -> Output: {'success' : False, 'info_value' : null}
Input: Running -> Output: {'success' : False, 'info_value' : null}
Input: Sarah -> Output: {'success' : True, 'info_value' : 'Sarah'}

The examples did not generalize perfectly, but they did improve the accuracy considerably.

Conclusion

All these changes improved the accuracy of the data collection. However, it was only 90% at best. Given I hope this to be a production ready system, this rate is not good enough. The next step will be to try a large model.

Leave a Reply