Vector Databases — Store and Search AI Embeddings

import numpy as np

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# Example with OpenAI embeddings
from openai import OpenAI
client = OpenAI()

def embed(text):
    resp = client.embeddings.create(
        input=text,
        model="text-embedding-3-small"
    )
    return np.array(resp.data[0].embedding)

vec_a = embed("machine learning model")
vec_b = embed("neural network training")
vec_c = embed("chocolate cake recipe")

print(cosine_similarity(vec_a, vec_b))  # ~0.87 (similar)
print(cosine_similarity(vec_a, vec_c))  # ~0.12 (different)

import faiss
import numpy as np

# Create 10,000 sample vectors of dimension 1536
dim = 1536
n_vectors = 10_000
vectors = np.random.randn(n_vectors, dim).astype('float32')

# Normalise for cosine similarity (FAISS uses inner product after normalisation)
faiss.normalize_L2(vectors)

# Build flat L2 index (exact, good for < 1M vectors)
index = faiss.IndexFlatIP(dim)   # Inner Product = cosine after normalisation
index.add(vectors)
print(f"Index size: {index.ntotal} vectors")

# Search — find 5 nearest neighbours to a query
query = np.random.randn(1, dim).astype('float32')
faiss.normalize_L2(query)

scores, indices = index.search(query, k=5)
print("Top 5 matches (indices):", indices[0])
print("Similarity scores:", scores[0])

# For large collections: use IVF (inverted file) index for speed
nlist = 100   # Number of clusters
quantizer = faiss.IndexFlatIP(dim)
index_ivf = faiss.IndexIVFFlat(quantizer, dim, nlist)
index_ivf.train(vectors)
index_ivf.add(vectors)
index_ivf.nprobe = 10  # Search 10 clusters (trade speed for recall)

import chromadb
from chromadb.utils import embedding_functions

# Persistent local database
client = chromadb.PersistentClient(path="./my_chroma_db")

# Use OpenAI embeddings automatically
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key="sk-...",
    model_name="text-embedding-3-small"
)

collection = client.get_or_create_collection(
    name="knowledge_base",
    embedding_function=openai_ef
)

# Add documents (ChromaDB embeds them automatically)
collection.add(
    documents=[
        "Python is a high-level programming language",
        "Machine learning requires large datasets",
        "The Eiffel Tower is in Paris, France",
    ],
    ids=["doc1", "doc2", "doc3"],
    metadatas=[{"source": "wiki"}, {"source": "textbook"}, {"source": "wiki"}]
)

# Query — semantic search
results = collection.query(
    query_texts=["coding with Python"],
    n_results=2,
    where={"source": "wiki"}   # Metadata filter
)
print(results["documents"])

from pinecone import Pinecone, ServerlessSpec

pc = Pinecone(api_key="your-api-key")

# Create index (one-time)
pc.create_index(
    name="knowledge-base",
    dimension=1536,
    metric="cosine",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)

index = pc.Index("knowledge-base")

# Upsert vectors with metadata
vectors = [
    ("id-001", [0.1, 0.2, ...], {"text": "Doc 1", "category": "tech"}),
    ("id-002", [0.3, 0.1, ...], {"text": "Doc 2", "category": "science"}),
]
index.upsert(vectors=vectors, namespace="production")

# Query
results = index.query(
    vector=[0.15, 0.18, ...],
    top_k=5,
    filter={"category": {"$eq": "tech"}},  # Metadata filtering
    include_metadata=True,
    namespace="production"
)
for match in results.matches:
    print(f"{match.id}: score={match.score:.3f}, text={match.metadata['text']}")