Vector stores

import hashlib
import re

import numpy as np

STOPWORDS = {"a", "an", "the", "is", "are", "be", "can", "do", "does", "i", "my", "to",
             "of", "in", "on", "for", "and", "or", "what", "how", "it", "you", "your", "at"}


def words(text: str) -> list[str]:
    """Lowercase, drop stopwords, and strip a few suffixes (refunded -> refund)."""
    out = []
    for w in re.findall(r"[a-z0-9-]+", text.lower()):
        if w in STOPWORDS:
            continue
        for suffix in ("ing", "ed", "s"):
            if w.endswith(suffix) and len(w) > len(suffix) + 3:
                w = w[: -len(suffix)]
                break
        out.append(w)
    return out


def embed(text: str, dim: int = 256) -> np.ndarray:
    """Toy embedding: hashed bag of words, L2-normalised. Real models capture meaning."""
    vec = np.zeros(dim, dtype=np.float32)
    for word in words(text):
        bucket = int(hashlib.md5(word.encode()).hexdigest(), 16) % dim
        vec[bucket] += 1.0
    norm = np.linalg.norm(vec)
    return vec / norm if norm else vec


class VectorStore:
    def __init__(self, dim: int = 256):
        self.dim = dim
        self.ids: list[str] = []
        self.records: list[dict] = []
        self.vectors = np.zeros((0, dim), dtype=np.float32)

    def add(self, records: list[dict]) -> None:
        new = np.stack([embed(r["text"], self.dim) for r in records])
        self.vectors = np.vstack([self.vectors, new])
        self.records.extend(records)
        self.ids.extend(r["id"] for r in records)

    def search(self, query: str, k: int = 3, where: dict | None = None) -> list[dict]:
        scores = self.vectors @ embed(query, self.dim)
        order = np.argsort(-scores)
        hits = []
        for i in order:
            rec = self.records[i]
            if where and any(rec.get(key) != value for key, value in where.items()):
                continue                                  # metadata filter
            hits.append({"id": rec["id"], "score": round(float(scores[i]), 3), "text": rec["text"]})
            if len(hits) == k:
                break
        return hits

    def delete(self, record_id: str) -> None:
        i = self.ids.index(record_id)
        self.vectors = np.delete(self.vectors, i, axis=0)
        del self.ids[i], self.records[i]

    def __len__(self) -> int:
        return len(self.ids)


store = VectorStore()
store.add([
    {"id": "kb-1", "text": "Annual plans can be refunded within 30 days.", "product": "vault", "lang": "en"},
    {"id": "kb-2", "text": "Backups run nightly and are kept for 90 days.", "product": "vault", "lang": "en"},
    {"id": "kb-3", "text": "Sync plans can be refunded within 14 days.", "product": "sync", "lang": "en"},
    {"id": "kb-4", "text": "Two-factor authentication is required for admins.", "product": "vault", "lang": "en"},
])

print(len(store), "records")
for hit in store.search("refund for an annual plan", k=2):
    print(hit)
print("filtered:", store.search("refund for an annual plan", k=1, where={"product": "sync"}))

store.delete("kb-3")
print(len(store), "records;", [h["id"] for h in store.search("refund", k=3)])

np.save("vectors.npy", store.vectors)                 # persist vectors
print(np.load("vectors.npy").shape)

import chromadb

client = chromadb.PersistentClient(path="./chroma_db")
collection = client.get_or_create_collection(name="kb", metadata={"hnsw:space": "cosine"})

collection.upsert(
    ids=["kb-1", "kb-2"],
    documents=["Annual plans can be refunded within 30 days.", "Backups run nightly and are kept for 90 days."],
    metadatas=[{"product": "vault"}, {"product": "vault"}],
)   # Chroma embeds documents with its default local model unless you pass embeddings=

results = collection.query(query_texts=["how do refunds work"], n_results=1, where={"product": "vault"})
print(results["ids"][0], results["documents"][0])

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE chunks (
    id          text PRIMARY KEY,
    source      text NOT NULL,
    product     text,
    content     text NOT NULL,
    embedding   vector(1536) NOT NULL,
    updated_at  timestamptz DEFAULT now()
);

CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops);

-- top 5 chunks for a query vector, filtered by product
-- (<=> is cosine distance, so 1 - distance = similarity)
SELECT id, source, content, 1 - (embedding <=> $1) AS similarity
FROM chunks
WHERE product = 'vault'
ORDER BY embedding <=> $1
LIMIT 5;

import faiss
import numpy as np

dim = 384
vectors = np.random.default_rng(0).random((10_000, dim), dtype=np.float32)
faiss.normalize_L2(vectors)                    # cosine via inner product

index = faiss.IndexHNSWFlat(dim, 32, faiss.METRIC_INNER_PRODUCT)
index.add(vectors)

query = vectors[:1].copy()
scores, ids = index.search(query, 3)
print("nearest id:", ids[0][0], "| similarity:", round(float(scores[0][0]), 3))