ChromaDB

Call ChromaDB from Flux tasks for embedding storage and similarity search.

Flux 0.56.0 does not ship a native ChromaDB integration. The pattern is to wrap a ChromaDB client in @task so embeddings and queries get the usual checkpoint, retry, and replay machinery on top of any persistent client.

Install

pip install chromadb

A persistent client writes to disk at a path you choose. The path needs to be reachable from every worker that runs the task — a shared volume in production, a local directory when developing inline.

A small workflow

import chromadb
from flux import workflow, task, ExecutionContext


def _client():
    return chromadb.PersistentClient(path="/var/lib/flux/chroma")


@task.with_options(cache=True)
async def embed_and_add(collection_name: str, docs: list[dict]) -> int:
    client = _client()
    coll = client.get_or_create_collection(collection_name)
    coll.upsert(
        ids=[d["id"] for d in docs],
        documents=[d["text"] for d in docs],
        metadatas=[d.get("metadata", {}) for d in docs],
    )
    return coll.count()


@task
async def query(collection_name: str, text: str, k: int = 5) -> list[dict]:
    client = _client()
    coll = client.get_collection(collection_name)
    result = coll.query(query_texts=[text], n_results=k)
    return [
        {"id": i, "doc": d, "distance": dist}
        for i, d, dist in zip(result["ids"][0], result["documents"][0], result["distances"][0])
    ]


@workflow
async def search(ctx: ExecutionContext[str]):
    await embed_and_add(
        "kb",
        [
            {"id": "a", "text": "Flux persists every state transition."},
            {"id": "b", "text": "Workers claim executions over SSE."},
            {"id": "c", "text": "ChromaDB is an embedding database."},
        ],
    )
    return await query("kb", ctx.input, k=2)

The cache=True flag on embed_and_add uses Flux’s disk-backed cache (flux/cache.py). The cache key is derived from the task source and arguments, so re-running the same workflow with the same inputs skips the embedding step entirely — useful when ChromaDB’s default sentence-transformer download is the slow part of your test loop. The cache persists across executions and across processes, so a worker restart does not invalidate it.

Embedding functions

Chroma’s default embedder is all-MiniLM-L6-v2, downloaded the first time the collection is created. To swap it (OpenAI, Cohere, Bedrock, your own), pass an embedding_function when calling get_or_create_collection. Pin the embedder per collection: a collection embedded with one model and queried with another returns nonsense, and Chroma will not warn you.

Multi-tenant collections

Chroma collections are flat. If you need per-tenant isolation, encode the tenant in the collection name (kb__acme, kb__widgetco) or in document metadata and filter at query time with where={"tenant": "acme"}. The same pattern works for environment separation (kb__staging, kb__prod).

What goes wrong

See also


Derived against ChromaDB 0.5.x, 2026-05.