Making AI Do the Math: A Simple Framework for Large Datasets

The Math Horror

Language Models by their nature are better with text than numbers.

Affinity Answers’ flagship AI application needs to compute metrics from a relatively large dataset of varying nature. Language Models by their nature are better with text than numbers. So, it was obvious that we need to relegate any arithmetic to a tool — say a calculator tool. Sounds simple, no? But to our horror with several thousand rows of datasets the game took a turn for the worse in terms of incorrectness of the results — forget latency and heavy token consumption for now. We developed a framework to handle math on a large dataset in AI applications by turning them into spreadsheet-like data for computations — which we fondly call as the Lotus-1–2–3 moment for our AI application. Read through the evolution of the framework from a simple calculator tool.

The Simple Calculator

One of the ways to make LLM better at Math is to give it access to a calculator tool, the calculator tool would accept mathematical expressions as an input and return the calculated result for the given mathematical expression as output.

Architecture with Calculator as a tool

For our initial implementation we took inspiration from Langchain Math. We improved on it by implementing better error handling for example, when LLM asks for something divided by zero instead of throwing an exception and breaking the agent we gently mention to the LLM that this division by zero cannot be computed so that LLM can handle it gracefully.

Limitations of the Simple Calculator

Limitations of Calculator with multiple calculations

The lacuna of this implementation as shown in the diagram above is that if the LLM has to perform many calculations it needs to perform tool calls multiple times, this increases the latency and cost.

Optimizing with Batched Calculator Calls

Batching tool calls by making calculator accept multiple expressions at once

We solved the above limitation by batching the multiple calculations needed to a single call as illustrated above. This was done by making the calculator tool accept multiple mathematical expressions at once, so the LLM can pass all the mathematical expressions it needs the answer to, reducing the number of tool calls, latency and cost.

Advanced agentic analytics tool with in-memory DB

LLMs performed poorly when the data it deals with is slightly larger than 100 records

While the calculator solved the problem of inaccuracies with mathematical expressions, there was an another inherent problem with LLM itself i.e. the problem with handling lot of data points at once, in our experimentations we noticed that LLMs performed poorly when the data it deals with is slightly larger than 100 records when we ask it to give me top 5 by a metric it sometimes used to skip certain records for the reason we do not understand, this lead us to believe for the scale we wanted to build an AI application ~1000’s of records it would not work well. The problem we faced could be related to the poor context efficiency.

The solution we thought for this was to give it Spreadsheet like capabilities so that it has few data transformation capabilities like

1. Sorting
2. Filterting
3. Computing Aggregations

While users often need a high-level summary of the data, our implementation also allows them to drill down into the specifics, especially after applying certain operations.

…we had a high level idea that giving it spreadsheet like capabilities can solve most of our analytical needs…

Now that we had a high level idea that giving it spreadsheet like capabilities can solve most of our analytical needs, the question arises that how do we implement this and make LLM control this efficiently, the answer was using a database that uses SQL, because SQL has all the capabilities that LLM needs and more importantly LLM can perform any operation it needs using a textual query.

For our use case specifically the data we deal with comes from various disparate sources such as different APIs due to operation complexities such as effort, data storage and to make sure the data LLM has access to is a realtime, it did not make sense to implement a CDC (Change data capture) pipeline for this. Instead we needed a real-time ad hoc data pipeline that can deal with just 10’s of thousands of data points. For this use case SQLite was perfect.

Now the question arises how do we create an Ad hoc data pipeline, the framework we came up with is as follows.

1. Transform the data from whatever format we have to a tabular format like CSV.
2. Select the fields necessary to answer the user’s query.
3. Put the data into an in memory SQLite db that expires when the current user request closes.
4. Give the LLM read-only access to this in memory db, so that it has data transformation capabilities.

Since the data comes from various disparate sources we cannot have a fixed schema, so the above framework is not deterministic and to handle the above requirement our Ad hoc data pipeline is LLM powered.

JMES Path is used to make the framework generic

Ad hoc LLM powered data pipeline to ETL JSON to In Memory DB

The implementation for our use case specifically will look like this, since the data we receive from the APIs is in JSON we had to use JMES Path to convert it to a tabular format. The JMES Path query itself is generated by an LLM that decides the fields necessary and the JMES path query needed to answer users’ query. Also in our implementation we do not pass the entire JSON array if the data is an array, instead we programmatically create a representative element that has all the fields and value types necessary for an LLM to generate the JMES Path reliably.

The final architecture with integration of Analytics tool

The final architecture will look like the following and with this our LLM can reliably perform data transformation operations and it’s faster and cost effective as well.

The positive side effects of this architecture other than obvious reliability are as follows.

1. We do not have to worry about data not fitting LLM’s context length anymore as data processing is done by SQLite.
2. It is faster and cheaper since not many prompt tokens are being used to send the data to LLM for processing instead LLM just has to generate SQL query to process the data.

The Lotus-1–2–3 moment

For our AI applications, which operate on volatile raw data streams (rather than pre-structured databases), this framework was our Lotus-1–2–3 moment. By transforming raw inputs into a spreadsheet-like computational layer, we enabled users to generate reliable, on-demand metrics from unbridled queries — without the instability and overhead that come from relying solely on the LLM.

Special thanks to Vivek Vijayan and Lakshay Batra both of whom have helped me in putting this blog together.

Was there a different way to solve this? Perhaps more elegant, a bit easier? Let us know in the comments.