Building Stateful Agent with Apache Pinot and Google Gen AI

In the era of Large Language Models (LLMs), building agents that can reason over data is a powerful paradigm. However, for an agent to be truly useful in a business context, it needs two critical capabilities: access to real-time analytical data and stateful memory of past interactions.

In this post, we explore how to build a Transaction Analyzer Agent that combines the low-latency analytical power of Apache Pinot with the reasoning capabilities of Google’s Gemini models.

The Goal

We will create an agent that can:

1. Analyze Transactions: Answer questions like “What are the last 10 transactions for the user X?” or “What is the average price of these orders?” by querying a Pinot database directly.
2. Maintain Context: Remember previous user queries and its own answers (e.g., “Can you provide more details about the retrieved transaction?”) to provide a natural, conversational experience.

The Stack

• Apache Pinot: Used as the storage engine for both the high-volume transaction data and the agent’s conversation memory.
• Google Gen AI: The brain of the agent, responsible for natural language understanding and tool orchestration.
• Python: The glue code that connects the LLM with Pinot using the pinotdb client.

By the end of this post, you’ll see how we architected this solution, set up the Pinot infrastructure, and implemented the agent’s memory persistence logic.

Architecture Overview

The system architecture is designed to be simple yet robust, separating the storage concerns from the reasoning logic.

Data Layer (Apache Pinot)

Apache Pinot serves as the backbone for our data needs. It hosts two distinct tables:

• **transaction**: A table containing the raw transaction records. This table is optimized for analytical queries, allowing the agent to perform aggregations (like calculating average price) or precise lookups in milliseconds.
• **memory**: A dedicated table for storing the conversation history. Each turn of the conversation (user inputs and agent responses) is ingested into this table, ordered by timestamp.

Application Layer

The agent logic is implemented in Python, leveraging the Google Gen AI SDK. It acts as the orchestrator:

• It is configured with a custom tool (query_pinot) that gives it permission to execute SQL queries against the Pinot Broker.
• It manages the conversation loop, ensuring that every interaction is recorded.

Interaction Flow

When a user interacts with the agent, the following sequence occurs:

1. User Query: The user sends a message (e.g., “Show me the last 5 orders”).
2. Context Retrieval: Before processing, the agent queries the memory table in Pinot to fetch the last 20 messages. This ensures the agent knows the current context.
3. Analysis & Tool Use: The Gemini model analyzes the user query + context. If it needs data, it calls the query_pinot tool to fetch transaction details from the transaction table.
4. Memory Persistence: The agent’s final response is generated. Crucially, both the user’s query and the agent’s response are immediately saved back to the memory table in Pinot.
5. Response: The final answer is delivered to the user.

Interaction Flow

This loop ensures that every subsequent query builds upon the history of the entire session.

Github repository: https://github.com/shrutimantri/stateful-agents-with-pinot

Pinot Infrastructure

A key part of making an agent intelligent is the infrastructure it relies on. For this project, we used Apache Pinot — a distributed OLAP data store — to handle both analytical data and conversational state.

Setting Up Pinot

At its core, a Pinot cluster consists of:

• Controller: Manages the cluster state and resource configuration.
• Broker: Receives queries from clients and scatters them across servers.
• Server: Hosts the data segments and executes queries.

In our setup, we used a local instance of Pinot 1.4.0. The initialization script ( init.py) ensures that the tables are clean and correctly configured before the agent starts.

Once the Pinot installation is complete, run the following commands that creates the vurtual environment, sets up env variables, and runs the init.py script:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
export PINOT_URL=http://localhost:9000
python init.py

Transaction Persistence

The transaction table stores historical e-commerce data. To make this data queryable by the agent, we defined a schema that treats the transaction time as a LONG epoch timestamp.

• Schema: Focuses on dimensions like userId, itemDescription, and country, with a ts column for time.
• Table Configuration: An OFFLINE table type designed for high-performance batch analytics.

Memory Persistence

The memory table is what makes the agent “stateful”. The memory data grows with every user interaction.

Schema:

• session_id: To uniquely identify different user conversations.
• role: Either user or model to distinguish who said what.
• content: The actual text of the message.
• ts: The exact time of the interaction.

Table Configuration: Also an OFFLINE table, we implemented a custom ingestion path to ensure new conversation turns are added to Pinot in near real-time.

By defining these two tables, we’ve given our agent a complete view of both the “World” (Transactions) and the “Conversation” (Memory).

Implementation Details

With our Pinot infrastructure in place, let’s dive into how we actually populated the tables and built the agent intelligence.

Data Ingestion and Date Parsing

One of the most common challenges in analytical data stores is handling diverse date formats. For our transaction data, we used a custom data generation script (datagen.py) that performs the following:

1. CSV Processing: Reads raw transaction data.
2. Date Normalization: Parses human-readable date strings (e.g., 12/01/2010 08:26) and converts them into milliseconds since epoch. This is crucial for Pinot's dateTimeFieldSpecs, allowing for efficient time-based filtering.
3. Segment Creation: Uses Pinot’s ingestion tools to convert this clean data into segments and push them to the controller.

Run the data generator using the following command:

python datagen/datagen.py

The Agent Brain (Gemini + Tool Use)

The core of our agent is built on Gemini 2.0 Flash. What makes it an “agent” rather than just a chatbot is its ability to use tools. We defined a Python function query_pinot() and exposed it to the model.

When the agent receives a request like “What are the last 5 orders for user X?”, the following happens:

• Reasoning: The model recognizes it needs specific data and decides to call query_pinot.
• Execution: It generates a precise SQL query (e.g., SELECT * FROM transaction WHERE userId = 'X' ORDER BY ts DESC LIMIT 5).
• Integration: The resulting data is returned to the model, which then summarizes it into a natural language response.

Managing Conversation State

To stay “in context,” the agent must remember what happened earlier in the conversation. We implemented this using a dual-write and single-read strategy:

1. Read (Context Retrieval): At the start of every message, the agent fetches the last 20 messages from the Pinot memory table based on the session_id.
2. Dual-Write (Persistence):

• Local Fallback: For immediate consistency, we write the latest messages to a local JSON file.
• Pinot Ingestion: Every response is also pushed to the Pinot memory table using the Controller’s /ingestFromFile REST endpoint. This ensures that even if our local script restarts, the conversation history is safely archived in our data store.

Demo & Results

To verify our stateful agent, we ran a series of conversational tests.

Query 1: Data Retrieval

The first step was to see if the agent could correctly use its tools to fetch data.

• User Question: “List the last 10 transactions for the userId 324429”
• Agent Logic: The agent queried the transaction table in Pinot, sorting by ts DESC.
• Result: It successfully listed the transaction IDs, items purchased, and timestamps for that user.

Query 2: Contextual Recall (The “Memory” Test)

This is where the stateful memory really shines. Instead of asking for data by ID again, we asked a follow-up question.

• User Question: “Can you just tell me the details of the second last order?”
• Agent Logic: The agent didn’t need to re-query the transaction table for all data. It looked into its conversation history (retrieved from the memory table), identified the orders listed in the previous turn, and accurately picked the “second last” one.
• Result: “The second last order has transactionId 6336605… and was made in the United Kingdom.”

Playing the demo

In order to play the demo yourself, run the following commands:

export GOOGLE_API_KEY=<YOUR_API_KEY>
export PINOT_HOST=localhost
export PINOT_BROKER_PORT=8000
export PINOT_URL=http://localhost:9000
python analyzer/agent.py

You can generate the Google API key from the AI Studio API Key’s page.

Conclusion

Building stateful agents requires more than just a smart LLM; it requires a robust data infrastructure. By using Apache Pinot as both an analytical engine and a conversational memory store, we achieved:

1. Low-Latency Analytics: Fast access to historical transaction data.
2. State Management: A persistent conversational state that survives application restarts.
3. Scalability: A platform capable of handling millions of transactions and thousands of concurrent agent sessions.

Future Work

While this demo uses Offline tables for simplicity, a production-grade implementation would leverage:

• Real-time Tables: To ingest memory via Kafka for lower latency.
• Upsert Feature: To allow updating transactions if needed.
• StarTree Cloud: For a managed experience in scaling Pinot clusters.

Hope this post gave you good insights into building stateful, data-driven agents!