Giving Agents Memory: Storage & Retrieval

An LLM has no memory of its own, every call starts from zero except for whatever text you put in the prompt. "Agent memory" means building a system that finds the right facts and hands them to the model at the right moment. This module covers both halves: where you store that information (a vector database), and how you make sure you retrieve the right pieces of it (retrieval quality).

What you'll learn

• Picking a vector database without overthinking the choice
• Why bigger chunks and smaller chunks each fail in different ways
• The one upgrade (reranking) that fixes most "wrong answer" retrieval bugs

1. The problem

Two different things go wrong with agent memory, and they need different fixes.

Storage: if you don't index your data for similarity search, "does this document relate to the user's question" means scanning everything, every time, slow, and it doesn't scale past a handful of documents.

Retrieval quality: even with a database in place, similarity search returns the closest matches, not necessarily the right ones. An agent that retrieves the wrong three paragraphs will confidently answer from the wrong context, and it won't tell you it did that. This is the failure mode that actually costs you: not "no results," but "confident, wrong results."

2. Options & when to use each

Where to store it

Option	Good for	Costs you	When to pick it
pgvector (Postgres)	Most projects, especially if you already use Postgres	Nothing extra to run, it's an extension	Default choice. You get SQL filtering (WHERE customer_id = ...) alongside similarity search for free.
A managed vector DB (Pinecone, etc.)	Teams that don't want to run their own database	Monthly cost, network round-trip per query	You're not already running Postgres and don't want to start
An embedded/local store (Chroma)	Prototypes, single-user tools, local-first apps	Doesn't scale well past one process	Fast to try something out; not for production with concurrent users

If you're already running Postgres for anything else in your stack, pgvector is almost always the right call, one less system to operate, and you can combine a normal WHERE clause with the similarity search (e.g., "similar documents, but only from this customer's account").

How to chunk and retrieve

Approach	Good for	Costs you	When to pick it
Straightforward chunking + top-k search	Getting started, most FAQs/docs use cases	Occasionally retrieves a chunk that's close but not quite right	Default starting point, build this first
Add a reranking step	Anywhere wrong answers are costly (support, legal, financial)	One extra model call per query (adds latency)	Your top-k search is finding relevant material but not always surfacing the best piece of it first

3. Build it

Store and search with pgvector

To initialize a vector search store in a PostgreSQL database, enable the vector extension and set up a document table with a vector column and HNSW index:

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    customer_id INT,
    embedding VECTOR(768)
);

CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops);

Next, write the Python integration code to generate embeddings using the Google GenAI SDK and run a cosine similarity query against the database table:

from google import genai

client = genai.Client()

def embed(text: str) -> list[float]:
    result = client.models.embed_content(
        model="gemini-embedding-2",
        contents=text,
        config={"output_dimensionality": 768},
    )
    return result.embeddings[0].values

def search(query: str, customer_id: int, top_k: int = 5):
    query_vec = embed(query)
    return db.execute(
        """
        SELECT content, 1 - (embedding <=> %s) AS similarity
        FROM documents
        WHERE customer_id = %s
        ORDER BY embedding <=> %s
        LIMIT %s
        """,
        [query_vec, customer_id, query_vec, top_k],
    )

The WHERE customer_id = %s alongside the similarity ordering is the practical reason to use Postgres over a pure vector store: you get relational filtering and semantic search in one query, so an agent can never accidentally retrieve another customer's data just because it's semantically similar.

Chunk documents sensibly

Don't index whole documents as single chunks (too big, you'll retrieve a 10-page PDF when the user asked about one paragraph) and don't index single sentences (too small, you lose surrounding context). A reasonable default:

def chunk_document(text: str, chunk_size: int = 800, overlap: int = 100) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap  # overlap keeps context from breaking mid-idea
    return chunks

This is a good starting point for most text. If you outgrow it, chunks routinely cut off mid-thought, or retrieval quality plateaus, the next step up is splitting on paragraph/section boundaries instead of a fixed character count, but don't add that complexity before you've confirmed you need it.

Add reranking when retrieval quality isn't good enough

Vector search alone finds the k closest matches fast, but "closest by embedding distance" and "actually most relevant to the question" aren't always the same thing. A reranker re-scores your top candidates using a model that looks at the query and each document together, not independently:

import cohere

co = cohere.Client()

def search_and_rerank(query: str, customer_id: int, top_k: int = 3):
    # Cast a slightly wider net with vector search first
    candidates = search(query, customer_id, top_k=15)

    reranked = co.rerank(
        query=query,
        documents=[c["content"] for c in candidates],
        top_n=top_k,
        model="rerank-english-v3.0",
    )
    return [candidates[r.index] for r in reranked.results]

Retrieve more candidates than you need (15, say), then let the reranker narrow it down to the best 3. This costs one extra API call per query, worth it anywhere a wrong answer is expensive, skippable for low-stakes internal tools.

4. What goes wrong

Mistake	How you notice it	The fix
No index on the embedding column	Queries get slower as the table grows	Add the hnsw index shown above, Postgres won't use it automatically until it exists
Chunks too large	Agent answers seem to pull in irrelevant tangents from the same document	Reduce chunk size; make sure overlap isn't zero
Chunks too small	Agent's answers feel fragmented, missing context that was right next to the retrieved piece	Increase chunk size or overlap
Filtering forgotten	One customer's agent can retrieve another customer's data	Always scope retrieval with a WHERE clause on tenant/customer ID, never rely on the embedding alone
Retrieval looks fine but answers are still wrong	Top result by similarity isn't actually the best answer to the question	Add a reranking step (above) , this is the single highest-leverage fix for "technically found something, but not the right thing"
Embedding model mismatch	vector dimensions do not match error, or silently bad results after switching models	Store which embedding model/version was used per row; re-embed everything if you switch models, old and new embeddings aren't comparable

5. Confirm it worked

To verify that your query execution plans are performing fast vector index scans rather than expensive full table scans, analyze your queries using EXPLAIN ANALYZE:

-- Confirm the index is actually being used, not a full table scan
EXPLAIN ANALYZE
SELECT content FROM documents
WHERE customer_id = 1
ORDER BY embedding <=> '[0.1, 0.2, ...]'
LIMIT 5;
-- Look for "Index Scan using documents_embedding_idx" in the output.
-- If you see "Seq Scan" instead, the index isn't being used.

Additionally, run a quick integration test in Python to ensure that queries return expected matches on basic semantic lookups:

# Confirm retrieval actually finds relevant content for a known question
results = search("What's your refund policy?", customer_id=1)
assert any("refund" in r["content"].lower() for r in results[:3])

If reranking is in place, spot-check a handful of real queries and confirm the top result after reranking is a better answer than the top result from vector search alone, that comparison is the actual test of whether reranking was worth adding.

Next Step: proceed to Long Context & Token Economics to manage context and cost as this agent's tasks grow.