Implementing memory with Firestore for long-running AI assistants

By Mariana Oliveira · 31 July 20267,434 views
Implementing memory with Firestore for long-running AI assistants

Introduction: The Memory Problem in Generative AI

When we build AI assistants today, we often fall into the trap of treating 'memory' as a simple flat array of strings passed to an LLM context window. In a real-world SaaS environment—like the one we are building here in Recife—that approach falls apart the moment a user transitions from a casual query to a long-running, collaborative workflow. As developers, we aren't just storing messages; we are building a state machine for an agent that needs to recall facts, summarize history, and maintain persona consistency across sessions that might last for months.

Firestore is uniquely suited for this, but its performance depends entirely on how you structure your memory persistence layer. If you ignore the underlying storage engine’s constraints, you’ll end up with massive documents that hit the 1MB limit or, worse, you'll burn through your read quotas by pulling entire conversation histories just to check one specific preference. In this article, we’ll explore the architecture of AI memory using the subcollection-versus-map framework.

The Decision Framework: Maps vs. Subcollections

Before you write a single line of code, you must decide how to store the conversation flow. This is the cornerstone of my decision framework.

Use an Embedded Map when:

  1. The data is small and always accessed together (e.g., the user’s 'current_persona' settings).
  2. You never need to query the internal items independently.
  3. The total size will consistently remain well under the 1MB document limit.

Use a Subcollection when:

  1. The data is unbounded or long-running (e.g., a message history that grows daily).
  2. You need to query specific parts of the data (e.g., 'give me all messages from yesterday' or 'find the last 10 messages').
  3. You need to leverage real-time listeners on the delta of new messages rather than the whole conversation history.

For most AI memory implementations, you will end up with a hybrid model: a parent thread document that stores metadata (like session start time and model parameters) and a messages subcollection that acts as the append-only log of the agent’s memory.

Designing the Schema for Scalability

Let’s look at a concrete implementation using Kotlin for our Android/Backend agent logic. In this structure, we treat the conversation thread as our root document and messages as the atomic unit of history.

// The Thread Metadata Document
data class ChatThread(
    val id: String,
    val userId: String,
    val createdAt: Timestamp,
    val modelVersion: String,
    val summary: String // Periodic summary of long-term memory
)

// The Message Subcollection Structure
data class Message(
    val role: String, // 'user' or 'assistant'
    val content: String,
    val timestamp: Timestamp,
    val tokens: Int,
    val metadata: Map<String, Any> // Optional context like tool outputs
)

Why this works: If we stored all messages in the ChatThread document as an array of maps, the document would quickly bloat. As the conversation length exceeds a few thousand tokens, you would hit the 1MB document limit. Furthermore, retrieving the 'last 10 messages' from a Firestore array requires pulling the entire array into memory and filtering it client-side. With a subcollection, we simply use a query with an orderBy and limit(10). This keeps your read costs predictable and your latency low.

Optimizing Real-time Listeners for Collaboration

Real-time collaborative AI requires that both the user and the agent see the same 'state of mind.' When a user is interacting with an AI, the agent is often processing a response in the background. Using an onSnapshot listener on the messages subcollection allows the UI to render the AI’s typing effect without a full page refresh.

However, a common pitfall is attaching a listener to the entire conversation history. If your thread is 5,000 messages long, the overhead of the listener will cause significant performance lag on mobile devices. Instead, apply the 'Windowing' technique.

  1. Query only the most recent N messages (e.g., last 20).
  2. Use a Firestore Query object:
// Kotlin example of a focused real-time listener
val messageQuery = firestore.collection("threads")
    .document(threadId)
    .collection("messages")
    .orderBy("timestamp", Query.Direction.DESCENDING)
    .limit(20)

messageQuery.addSnapshotListener { snapshots, e ->
    // Handle real-time updates for only the active window
    // This prevents massive bandwidth consumption
}

This pattern ensures your AI assistant remains snappy. By limiting the scope of the real-time listener to the 'active context window,' you keep the client-side state synchronized without overwhelming the device memory or Firestore’s read-per-document cost.

Consistency and The 'Long-Term' Memory Bridge

While subcollections are perfect for current conversations, what about 'long-term memory'? If your AI needs to recall something the user said three weeks ago, you shouldn't be querying the entire subcollection history. This is where you introduce a 'Memory Index.'

Every time a conversation reaches a certain token threshold, trigger a Cloud Function to summarize the context into a 'Memory Document.'

  1. The Current Subcollection: Holds the raw, real-time message stream.
  2. The Summary Field: Lives on the parent Thread document. It is updated periodically by an LLM-based summarizer.
  3. The Global Memory Store: A top-level collection where vectors (embeddings) are stored, allowing for semantic search (RAG) across all conversations.

This tiered approach is crucial. The user doesn't need to see the entire history in the UI chat window; they only need the current context. The LLM doesn't need to see every message ever sent; it needs the active context plus the summarized long-term memory. This architecture prevents the 'context window exhaustion' that plagues naive AI implementations.

Handling Concurrent Writes in AI Workflows

In a collaborative app where multiple users might contribute to a thread—or where an agent is running multiple background tool-use processes—you must worry about write contention. Firestore handles this with ACID transactions, but you should avoid them for high-frequency chat updates if possible.

Instead, use the serverTimestamp and local optimistic updates. If you allow multiple participants to 'inject' memory into the thread, treat the subcollection as an append-only log. Since Firestore documents have a strict write limit of 1 write per second, ensure your AI agent is not hammering the same document. By keeping the messages as individual documents in a subcollection, you avoid the contention on the parent document. Each new message is a separate write operation, allowing your AI to handle high-velocity inputs without hitting document write throughput bottlenecks.

Performance Trade-offs and Best Practices

Let’s summarize the trade-offs we have navigated:

  • Read-Heavy Patterns: If your users frequently look back at their history, the subcollection approach is superior because you can use indexes to fetch specific slices of data. If you used a single giant array in a map, your read cost would scale linearly with the total conversation size, making the app feel sluggish over time.
  • Write-Heavy Patterns: AI agents can be chatty. If an agent emits a response token-by-token (streaming), do not try to write every token to Firestore. Aggregate the stream into chunks of 100-200 tokens before committing a write to the messages subcollection. This keeps your billing manageable and your write-rate usage within healthy bounds.
  • Storage Limits: Always remember the 1MB document limit. By choosing subcollections, you effectively remove the ceiling on your conversation length. You could technically store millions of messages in a single thread without ever breaking the underlying database constraints, provided your queries are properly indexed.

Final Thoughts: The Path Forward

Building an AI assistant that remembers is not just about the prompt engineering; it is about how you treat the data lifecycle. In our studio, we’ve found that the simplest architecture is often the most resilient. Start with a clear separation between your metadata (the Thread document) and your transient state (the Messages subcollection).

As your startup grows, you might eventually need to move your vector-based long-term memory into a dedicated engine like Pinecone, while keeping your short-term state in Firestore. That is a natural evolution, and by following the framework I’ve outlined, you will have the flexibility to switch out your backend components without needing to rewrite your entire data access layer.

Remember, your users value consistency above all else. They want an AI that 'knows' them. By modeling your data to support efficient retrieval and real-time synchronization, you are building the foundation of a true digital companion rather than just a chatbot. Start small, use subcollections for your history, keep your listeners scoped to the active window, and scale your persistence as your AI’s 'brain' grows.

Firestore gives you the tools to handle the heavy lifting, but the architecture is up to you. Treat the data as a growing tree, not a static snapshot, and you’ll find that building for AI becomes significantly more predictable and, ultimately, much more rewarding.

Comments

No comments yet. Be the first!

Sign in to leave a comment.