Give your AI a memory on your own disk with Chroma

What you get

A Python script that saves a few short notes into a Chroma database folder on your disk and finds the ones closest in meaning to a question, with no cloud database and no API key.

Prerequisites

Python 3.9 or newer (the chromadb package requires >=3.9)
pip
An internet connection for the first run. Chroma's default embedding function uses the all-MiniLM-L6-v2 model, runs on your machine, and may download the model files automatically the first time. After that, nothing leaves your machine.

Steps

Install the official Python package.

Shell

pip install chromadb==1.5.9

Create a file named memory.py with this code.

Python

import chromadb

# Stores the database files in this folder and loads them on the next run.
client = chromadb.PersistentClient(path="./chroma_data")

# No embedding function is given, so Chroma uses its default one.
collection = client.get_or_create_collection(name="notes")

# A record whose id already exists is ignored, so re-running is safe.
collection.add(
    ids=["note1", "note2", "note3"],
    documents=[
        "The office wifi password is changed on the first Monday of each month.",
        "Maria prefers meetings in the afternoon, never before 10am.",
        "The quarterly report is due on the last Friday of March.",
    ],
)

results = collection.query(
    query_texts=["When does Maria like to meet?"],
    n_results=2,
)

for doc, distance in zip(results["documents"][0], results["distances"][0]):
    print(round(distance, 3), doc)

Know what each part does.

PersistentClient(path=...) saves data automatically to that folder and loads it on start if it exists.
get_or_create_collection creates the collection the first time and returns the existing one after that.
add takes your text and a unique string id for each record. Chroma embeds and indexes the text for you.
query embeds the question and returns the n_results closest documents. Results are nested lists, one inner list per question, which is why the code reads [0].

Test it

Run the script:

Shell

python memory.py

The first run may pause while the default model files download. You should then see two lines, each a distance number followed by a note. The note about Maria's meetings should come first, with the smallest distance, even though the question does not use the same words. Exact distance numbers vary.

A new chroma_data folder appears next to the script. Run the script again: it prints the same two notes, now read from that folder, and adds no duplicates.