Vector databases: pgvector, Chroma and friends

pip install -U chromadb

import chromadb

from embed_util import embed

# A persistent client stores data in this folder, so it survives restarts
client = chromadb.PersistentClient(path="chroma_data")

collection = client.get_or_create_collection(
    name="menu",
    metadata={"hnsw:space": "cosine"},  # use cosine distance
)

ITEMS = [
    {"id": "espresso", "text": "Espresso: a short, intense shot of coffee.", "category": "coffee"},
    {"id": "cappuccino", "text": "Cappuccino: espresso with steamed milk and a thick layer of foam.", "category": "coffee"},
    {"id": "cold-brew", "text": "Cold Brew: coffee steeped slowly in cold water, smooth, bold and unsweetened.", "category": "coffee"},
    {"id": "filter-coffee", "text": "Filter Coffee: South Indian style coffee with hot milk, served in a steel tumbler.", "category": "coffee"},
    {"id": "masala-chai", "text": "Masala Chai: black tea simmered with milk and spices like ginger and cardamom.", "category": "tea"},
    {"id": "croissant", "text": "Butter Croissant: flaky, buttery pastry baked fresh every morning.", "category": "bakery"},
    {"id": "banana-bread", "text": "Banana Bread: moist, lightly sweet loaf made with ripe bananas.", "category": "bakery"},
    {"id": "veg-sandwich", "text": "Veg Grilled Sandwich: grilled bread with vegetables and cheese.", "category": "food"},
]

vectors = embed([i["text"] for i in ITEMS], "RETRIEVAL_DOCUMENT")

# upsert = insert, or replace if the id already exists (safe to re-run)
collection.upsert(
    ids=[i["id"] for i in ITEMS],
    documents=[i["text"] for i in ITEMS],
    embeddings=vectors.tolist(),
    metadatas=[{"category": i["category"]} for i in ITEMS],
)
print("Stored:", collection.count())

import chromadb

from embed_util import embed

collection = chromadb.PersistentClient(path="chroma_data").get_collection("menu")


def search(query, k=3, category=None):
    q = embed([query], "RETRIEVAL_QUERY")[0].tolist()
    result = collection.query(
        query_embeddings=[q],
        n_results=k,
        where={"category": category} if category else None,  # metadata filter
    )
    # Results are lists of lists: one inner list per query
    for doc, dist in zip(result["documents"][0], result["distances"][0]):
        print(f"  similarity {1 - dist:.2f}  {doc}")


print("Anything warm and spicy?")
search("anything warm and spicy?")

print("\nSomething sweet, bakery only")
search("something sweet", category="bakery")

docker run --name pgvec -e POSTGRES_PASSWORD=postgres -p 5432:5432 -d pgvector/pgvector:pg17
pip install -U "psycopg[binary]" pgvector

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS menu_docs (
    id        text PRIMARY KEY,
    content   text NOT NULL,
    category  text NOT NULL,
    embedding vector(768) NOT NULL   -- must match the embedding size you use
);

-- Approximate index for fast cosine search on large tables.
-- With only 8 rows Postgres will not use it, but this is what you add at scale.
CREATE INDEX IF NOT EXISTS menu_docs_embedding_idx
    ON menu_docs USING hnsw (embedding vector_cosine_ops);

import psycopg
from pgvector.psycopg import register_vector

from chroma_index import ITEMS  # reuse the same menu list (this re-runs chroma_index.py, which is harmless)
from embed_util import embed

DSN = "postgresql://postgres:postgres@localhost:5432/postgres"

with psycopg.connect(DSN, autocommit=True) as conn:
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")  # before register_vector
    register_vector(conn)  # lets psycopg send and read NumPy vectors

    conn.execute(open("schema.sql", encoding="utf-8").read())

    vectors = embed([i["text"] for i in ITEMS], "RETRIEVAL_DOCUMENT")
    for item, vec in zip(ITEMS, vectors):
        conn.execute(
            """INSERT INTO menu_docs (id, content, category, embedding)
               VALUES (%s, %s, %s, %s)
               ON CONFLICT (id) DO UPDATE
               SET content = EXCLUDED.content,
                   category = EXCLUDED.category,
                   embedding = EXCLUDED.embedding""",
            (item["id"], item["text"], item["category"], vec),
        )
    print("Rows:", conn.execute("SELECT count(*) FROM menu_docs").fetchone()[0])

import psycopg
from pgvector.psycopg import register_vector

from embed_util import embed

DSN = "postgresql://postgres:postgres@localhost:5432/postgres"


def search(query, k=3, category=None):
    q = embed([query], "RETRIEVAL_QUERY")[0]
    with psycopg.connect(DSN) as conn:
        register_vector(conn)
        rows = conn.execute(
            """SELECT content, 1 - (embedding <=> %s) AS similarity
               FROM menu_docs
               WHERE (%s::text IS NULL OR category = %s)
               ORDER BY embedding <=> %s
               LIMIT %s""",
            (q, category, category, q, k),
        ).fetchall()
    for content, similarity in rows:
        print(f"  similarity {similarity:.2f}  {content}")


if __name__ == "__main__":
    search("anything warm and spicy?")
    search("something sweet", category="bakery")

I am building [describe app, e.g. a college notes search / support bot]. Data: about [number] documents, [number] users, hosted on [cloud or laptop], team skills: [skills]. Existing database: [Postgres / MySQL / none]. Recommend Chroma, pgvector, a dedicated vector database or a managed service. Give the reasons, the risks and what I should measure before committing.