Vector Databases & AI Integration

# Install: pip install openai

from openai import OpenAI

client = OpenAI(api_key="sk-...")  # or read OPENAI_API_KEY from the environment

# Generate text embedding
response = client.embeddings.create(
    model="text-embedding-3-small",  # 1536 dimensions
    input="The quick brown fox jumps over the lazy dog"
)

embedding = response.data[0].embedding
print(f"Embedding dimensions: {len(embedding)}")
print(f"First 5 values: {embedding[:5]}")

# Pinecone Local is the official in-memory emulator: the same API and the
# same client SDK, no account and no API key. It stores nothing on disk,
# so every restart gives you an empty instance.
docker run -d --name pinecone-local \
  -e PORT=5080 -e PINECONE_HOST=localhost \
  -p 5080-5090:5080-5090 --platform linux/amd64 \
  ghcr.io/pinecone-io/pinecone-local:latest

# 5080 is the control plane (create/describe/delete index). Each index you
# create gets its own port from the mapped range, starting at 5081, which
# is why the range and not just one port has to be published.
curl -s http://localhost:5080/indexes

# Finished? Remove it, it holds no state you need:
#   docker rm -f pinecone-local

python3 -m venv .venv && source .venv/bin/activate
pip install pinecone sentence-transformers

from pinecone import Pinecone, ServerlessSpec

# Pinecone Local ignores the key, but the SDK insists one is present.
# Against the real service you would drop host= and pass your own key.
pc = Pinecone(api_key="pclocal", host="http://localhost:5080")

pc.create_index(
    name="semantic-search",
    dimension=384,    # must match the embedding model exactly
    metric="cosine",  # cosine, euclidean or dotproduct
    spec=ServerlessSpec(cloud="aws", region="us-east-1"),
)

desc = pc.describe_index("semantic-search")
print(f"name={desc.name} dimension={desc.dimension} metric={desc.metric}")
print(f"host={desc.host} ready={desc.status.ready}")

from pinecone import Pinecone
from sentence_transformers import SentenceTransformer

# A real embedding model, 384 dimensions, runs locally with no API key.
# Swap this for the OpenAI call shown earlier and the index becomes 1536-dim.
model = SentenceTransformer("all-MiniLM-L6-v2")

pc = Pinecone(api_key="pclocal", host="http://localhost:5080")

# Pinecone Local advertises an https:// host but serves plain HTTP.
index_host = pc.describe_index("semantic-search").host.replace("https://", "http://")
index = pc.Index(host=index_host)

docs = [
    ("doc1", "Python is a programming language",   "tech"),
    ("doc2", "Dogs are loyal companions",          "animals"),
    ("doc3", "JavaScript powers web applications", "tech"),
]

vectors = [
    {
        "id": doc_id,
        "values": model.encode(text).tolist(),
        "metadata": {"text": text, "category": category},
    }
    for doc_id, text, category in docs
]

# Upsert = insert or overwrite by id. Re-running this is safe.
index.upsert(vectors=vectors)
print(f"Upserted {len(vectors)} vectors")
print(index.describe_index_stats())

from pinecone import Pinecone
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
pc = Pinecone(api_key="pclocal", host="http://localhost:5080")
index_host = pc.describe_index("semantic-search").host.replace("https://", "http://")
index = pc.Index(host=index_host)

# Note: not one word of this query appears in any stored document.
query_text = "What are the best coding languages?"
query_embedding = model.encode(query_text).tolist()

results = index.query(vector=query_embedding, top_k=3, include_metadata=True)

for match in results["matches"]:
    print(f"Score: {match['score']:.4f}")
    print(f"Text: {match['metadata']['text']}")
    print(f"Category: {match['metadata']['category']}")
    print()

from pinecone import Pinecone
from sentence_transformers import SentenceTransformer


model = SentenceTransformer("all-MiniLM-L6-v2")
pc = Pinecone(api_key="pclocal", host="http://localhost:5080")
index_host = pc.describe_index("semantic-search").host.replace("https://", "http://")
index = pc.Index(host=index_host)
query_embedding = model.encode("What are the best coding languages?").tolist()

# The filter is applied during the search, not after it: asking for
# top_k=10 still returns only vectors whose metadata matches.
results = index.query(
    vector=query_embedding,
    top_k=10,
    filter={"category": {"$eq": "tech"}},
    include_metadata=True,
)

for match in results["matches"]:
    print(f"{match['score']:.4f}: {match['metadata']['text']}")

# Milvus normally runs as three services (Milvus, etcd for metadata and
# MinIO for object storage). The standalone image can embed etcd and use
# local disk instead, which collapses the whole thing into one container.
docker run -d --name milvus-standalone \
  -p 19530:19530 -p 9091:9091 \
  -e ETCD_USE_EMBED=true \
  -e ETCD_DATA_DIR=/var/lib/milvus/etcd \
  -e COMMON_STORAGETYPE=local \
  milvusdb/milvus:v2.5.4 milvus run standalone

# 19530 is the gRPC port the SDK uses; 9091 serves health and metrics.
# Startup takes ~30s, and the client will refuse connections until then.
curl -s http://localhost:9091/healthz

# Finished? Remove it, it holds no state you need:
#   docker rm -f milvus-standalone

python3 -m venv .venv && source .venv/bin/activate
pip install pymilvus sentence-transformers

from pymilvus import MilvusClient, DataType

client = MilvusClient(uri="http://localhost:19530")

if client.has_collection("documents"):
    client.drop_collection("documents")

# auto_id lets Milvus assign the primary key on insert.
schema = MilvusClient.create_schema(auto_id=True, enable_dynamic_field=False)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=384)
schema.add_field("text", DataType.VARCHAR, max_length=1000)
schema.add_field("category", DataType.VARCHAR, max_length=50)

client.create_collection(collection_name="documents", schema=schema)

info = client.describe_collection("documents")
print(f"Collection: {info['collection_name']}")
print(f"Fields: {[f['name'] for f in info['fields']]}")

from pymilvus import MilvusClient, DataType

client = MilvusClient(uri="http://localhost:19530")

index_params = client.prepare_index_params()
index_params.add_index(
    field_name="embedding",
    index_type="HNSW",        # Hierarchical Navigable Small World
    metric_type="COSINE",
    params={
        "M": 16,              # edges kept per node, per layer
        "efConstruction": 256,  # candidate list size while building
    },
)

client.create_index(collection_name="documents", index_params=index_params)

# A collection must be loaded into memory before it can be searched.
client.load_collection("documents")

for idx in client.list_indexes("documents"):
    d = client.describe_index("documents", idx)
    print(f"index={idx} type={d['index_type']} metric={d['metric_type']}")
print(f"load state: {client.get_load_state('documents')['state']}")

from pymilvus import MilvusClient
from sentence_transformers import SentenceTransformer


client = MilvusClient(uri="http://localhost:19530")
model = SentenceTransformer("all-MiniLM-L6-v2")

docs = [
    ("Machine learning models require training data", "AI"),
    ("Neural networks use backpropagation",           "AI"),
    ("Databases store structured information",        "databases"),
]

# Row-oriented dicts, one per entity. "id" is omitted because auto_id=True.
data = [
    {"embedding": model.encode(text).tolist(), "text": text, "category": category}
    for text, category in docs
]

result = client.insert(collection_name="documents", data=data)
print(f"Inserted {result['insert_count']} documents")

# Milvus writes are not searchable until they are sealed and indexed.
# flush() forces that now, which is what makes this example deterministic.
client.flush("documents")
print(client.get_collection_stats("documents"))

from pymilvus import MilvusClient
from sentence_transformers import SentenceTransformer


client = MilvusClient(uri="http://localhost:19530")
model = SentenceTransformer("all-MiniLM-L6-v2")
query_embedding = model.encode("How do models learn from examples?").tolist()

results = client.search(
    collection_name="documents",
    data=[query_embedding],
    limit=3,
    filter='category == "AI"',   # Milvus expression language, not JSON
    output_fields=["text", "category"],
    search_params={
        "metric_type": "COSINE",
        "params": {"ef": 64},    # candidates held during the walk
    },
)

for hits in results:
    for hit in hits:
        print(f"Score: {hit['distance']:.4f}")
        print(f"Text: {hit['entity']['text']}")
        print(f"Category: {hit['entity']['category']}")

# Weaviate's selling point is that it vectorizes for you, which means it
# needs a model to do it with. Point it at an inference container running
# the same all-MiniLM-L6-v2 used earlier, and no API key is involved.
docker network create weaviate-net

docker run -d --name t2v-transformers --network weaviate-net \
  -e ENABLE_CUDA=0 \
  cr.weaviate.io/semitechnologies/transformers-inference:sentence-transformers-all-MiniLM-L6-v2

docker run -d --name weaviate --network weaviate-net \
  -p 8080:8080 -p 50051:50051 \
  -e ENABLE_MODULES=text2vec-transformers \
  -e DEFAULT_VECTORIZER_MODULE=text2vec-transformers \
  -e TRANSFORMERS_INFERENCE_API=http://t2v-transformers:8080 \
  -e AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED=true \
  -e PERSISTENCE_DATA_PATH=/var/lib/weaviate \
  cr.weaviate.io/semitechnologies/weaviate:1.34.1

# 8080 is REST, 50051 is gRPC. The v4 client needs BOTH: it does schema
# work over REST and batching and search over gRPC, so publishing only
# 8080 gets you a client that connects and then times out on insert.
curl -s -o /dev/null -w "%{http_code}\n" http://localhost:8080/v1/.well-known/ready

# Finished? Remove both, plus the network:
#   docker rm -f weaviate t2v-transformers && docker network rm weaviate-net

python3 -m venv .venv && source .venv/bin/activate
pip install weaviate-client

import weaviate
from weaviate.classes.config import Configure, DataType, Property

client = weaviate.connect_to_local()

if client.collections.exists("Article"):
    client.collections.delete("Article")

client.collections.create(
    name="Article",
    description="A blog article",
    # Weaviate vectorizes on write: hand it text, not embeddings.
    vector_config=Configure.Vectors.text2vec_transformers(),
    properties=[
        Property(name="title",       data_type=DataType.TEXT),
        Property(name="content",     data_type=DataType.TEXT),
        Property(name="category",    data_type=DataType.TEXT),
        Property(name="publishDate", data_type=DataType.DATE),
    ],
)

config = client.collections.use("Article").config.get()
print(f"Collection: {config.name}")
print(f"Properties: {[p.name for p in config.properties]}")
client.close()

from datetime import datetime, timezone

import weaviate

client = weaviate.connect_to_local()
articles = client.collections.use("Article")

rows = [
    {
        "title": "Introduction to Python",
        "content": "Python is a versatile programming language",
        "category": "Programming",
        "publishDate": datetime(2025, 1, 15, tzinfo=timezone.utc),
    },
    {
        "title": "Deep Learning Basics",
        "content": "Neural networks are the foundation of deep learning",
        "category": "AI",
        "publishDate": datetime(2025, 2, 1, tzinfo=timezone.utc),
    },
    {
        "title": "Database Design Patterns",
        "content": "Normalization and indexing are key to performance",
        "category": "Databases",
        "publishDate": datetime(2025, 1, 20, tzinfo=timezone.utc),
    },
]

# No embeddings anywhere in this payload: the vectorizer produces them
# server side, from the text properties, as each object is written.
with articles.batch.fixed_size(batch_size=100) as batch:
    for row in rows:
        batch.add_object(properties=row)

# Batch errors do not raise. They are collected, and ignoring them is
# how you end up with a silently half-loaded collection.
if articles.batch.failed_objects:
    raise RuntimeError(articles.batch.failed_objects)

print(f"Inserted {len(rows)} articles")
print(f"Objects in collection: {articles.aggregate.over_all(total_count=True).total_count}")
client.close()

import weaviate
from weaviate.classes.query import MetadataQuery

client = weaviate.connect_to_local()
articles = client.collections.use("Article")

# near_text, not near_vector: the query string is vectorized server side
# by the same module that vectorized the objects.
response = articles.query.near_text(
    query="machine learning and artificial intelligence",
    limit=2,
    return_properties=["title", "content", "category"],
    return_metadata=MetadataQuery(distance=True),
)

for obj in response.objects:
    print(f"Distance: {obj.metadata.distance:.4f}")
    print(f"Title: {obj.properties['title']}")
    print(f"Category: {obj.properties['category']}")
    print()

client.close()

import weaviate
from weaviate.classes.query import Filter, MetadataQuery

client = weaviate.connect_to_local()
articles = client.collections.use("Article")

# alpha=0.5 weights the two rankings evenly: 0 is pure keyword (BM25),
# 1 is pure vector. Weaviate fuses the two result lists into one score.
response = articles.query.hybrid(
    query="neural networks",
    alpha=0.5,
    filters=Filter.by_property("category").equal("AI"),
    limit=3,
    return_properties=["title", "category"],
    return_metadata=MetadataQuery(score=True),
)

for obj in response.objects:
    print(f"{obj.metadata.score:.4f}  {obj.properties['title']} ({obj.properties['category']})")

client.close()

import numpy as np
from numpy.linalg import norm

# Sample vectors
vec1 = np.array([1.0, 2.0, 3.0])
vec2 = np.array([1.5, 2.2, 2.8])

# Cosine similarity
cosine = np.dot(vec1, vec2) / (norm(vec1) * norm(vec2))
print(f"Cosine similarity: {cosine:.4f}")

# Euclidean distance
euclidean = norm(vec1 - vec2)
print(f"Euclidean distance: {euclidean:.4f}")

# Dot product
dot_prod = np.dot(vec1, vec2)
print(f"Dot product: {dot_prod:.4f}")

from pinecone import Pinecone, ServerlessSpec
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
pc = Pinecone(api_key="pclocal", host="http://localhost:5080")

# RAG needs a corpus with something to say. The three toy sentences from
# the Pinecone walkthrough are too thin: ask them a real question and a
# well-behaved model correctly answers "the context does not contain it".
KNOWLEDGE = [
    "Vector databases store embeddings and retrieve them by similarity rather than by exact match.",
    "An approximate index such as HNSW trades a small amount of recall for a very large speed gain.",
    "Without a vector index every query must compare the query against all stored vectors.",
    "Metadata filters narrow a similarity search to a subset, such as one tenant or one category.",
    "Embeddings from different models are not comparable, so changing models means reindexing.",
    "Recommendation systems use the same nearest-neighbour search as semantic search.",
]

if not pc.has_index("rag-docs"):
    pc.create_index(
        name="rag-docs",
        dimension=384,
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
    )

index_host = pc.describe_index("rag-docs").host.replace("https://", "http://")
index = pc.Index(host=index_host)

index.upsert(vectors=[
    {"id": str(i), "values": model.encode(text).tolist(), "metadata": {"text": text}}
    for i, text in enumerate(KNOWLEDGE)
])
print(f"Indexed {len(KNOWLEDGE)} chunks")


# The query MUST be embedded with the model that embedded the documents.
def retrieve_context(query, top_k=3):
    query_embedding = model.encode(query).tolist()
    results = index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
    return [match["metadata"]["text"] for match in results["matches"]]


question = "Why does changing the embedding model force a reindex?"
print(f"\nQuestion: {question}\n")
print("Retrieved context:")
for chunk in retrieve_context(question):
    print(f"  - {chunk}")

from openai import OpenAI
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer

# Runs standalone. It reads the rag-docs index the previous block created,
# so run that one first, but nothing here depends on its variables.
model = SentenceTransformer("all-MiniLM-L6-v2")
pc = Pinecone(api_key="pclocal", host="http://localhost:5080")
index_host = pc.describe_index("rag-docs").host.replace("https://", "http://")
index = pc.Index(host=index_host)

# The only piece that needs a real key. Retrieval needed nothing.
client = OpenAI(api_key="sk-...")


def retrieve_context(query, top_k=3):
    query_embedding = model.encode(query).tolist()
    results = index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
    return [match["metadata"]["text"] for match in results["matches"]]


def answer_question(question):
    context = "\n\n".join(retrieve_context(question))

    prompt = f"""Answer the question using only the context below.
If the context does not contain the answer, say so.

Context:
{context}

Question: {question}

Answer:"""

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content


print(answer_question("Why does changing the embedding model force a reindex?"))

# Pinecone Local and Milvus are already running from the sections above.
# The two services this section adds are Kafka and Spark.

# Kafka in KRaft mode: one container, no ZooKeeper.
docker run -d --name kafka-demo -p 29092:29092 \
  -e KAFKA_NODE_ID=1 \
  -e KAFKA_PROCESS_ROLES=broker,controller \
  -e KAFKA_LISTENERS=PLAINTEXT://0.0.0.0:29092,CONTROLLER://0.0.0.0:9093 \
  -e KAFKA_ADVERTISED_LISTENERS=PLAINTEXT://localhost:29092 \
  -e KAFKA_CONTROLLER_LISTENER_NAMES=CONTROLLER \
  -e KAFKA_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT \
  -e KAFKA_CONTROLLER_QUORUM_VOTERS=1@localhost:9093 \
  -e KAFKA_INTER_BROKER_LISTENER_NAME=PLAINTEXT \
  -e KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR=1 \
  -e KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR=1 \
  -e KAFKA_TRANSACTION_STATE_LOG_MIN_ISR=1 \
  -e CLUSTER_ID=5L6g3nShT-eMCtK--X86sw \
  apache/kafka:3.8.0

docker exec kafka-demo /opt/kafka/bin/kafka-topics.sh \
  --bootstrap-server localhost:29092 \
  --create --topic documents --partitions 1 --replication-factor 1

# Spark needs no container: PySpark runs a local cluster in-process, and
# it needs a JDK (17 or 21) on PATH.
python3 -m venv .venv && source .venv/bin/activate
pip install kafka-python-ng pyspark pinecone pymilvus sentence-transformers

# Finished? Remove it:
#   docker rm -f kafka-demo

import json

from kafka import KafkaProducer

producer = KafkaProducer(
    bootstrap_servers=["localhost:29092"],
    value_serializer=lambda v: json.dumps(v).encode("utf-8"),
)

DOCS = [
    {"id": "1", "title": "Sharding", "source": "wiki",
     "content": "Sharding splits one logical table across many physical machines."},
    {"id": "2", "title": "Replication", "source": "wiki",
     "content": "Replication keeps copies of the same data on more than one node."},
    {"id": "3", "title": "Partitioning", "source": "wiki",
     "content": "Partitioning divides a table into smaller pieces inside one database."},
]

for doc in DOCS:
    producer.send("documents", doc)
producer.flush()
print(f"Produced {len(DOCS)} documents to topic 'documents'")

import json

from kafka import KafkaConsumer
from pinecone import Pinecone, ServerlessSpec
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# The index dimension must match the embedding model. 384 here because
# that is what all-MiniLM-L6-v2 produces; an index built at 1536 for
# OpenAI would reject every one of these upserts.
pc = Pinecone(api_key="pclocal", host="http://localhost:5080")
if not pc.has_index("documents"):
    pc.create_index(
        name="documents",
        dimension=384,
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
    )
index_host = pc.describe_index("documents").host.replace("https://", "http://")
index = pc.Index(host=index_host)

consumer = KafkaConsumer(
    "documents",
    bootstrap_servers=["localhost:29092"],
    auto_offset_reset="earliest",
    # Without this the loop blocks forever, which is correct in production
    # and unhelpful in a demo. It stops after 10s with no new messages.
    consumer_timeout_ms=10000,
    value_deserializer=lambda m: json.loads(m.decode("utf-8")),
)

count = 0
for message in consumer:
    doc = message.value
    embedding = model.encode(doc["content"]).tolist()

    index.upsert(vectors=[{
        "id": doc["id"],
        "values": embedding,
        "metadata": {"title": doc["title"], "source": doc["source"]},
    }])
    count += 1
    print(f"Indexed document {doc['id']}: {doc['title']}")

print(f"\nConsumed and indexed {count} documents")
print(index.describe_index_stats())

from pymilvus import MilvusClient, DataType
from pyspark.sql import SparkSession

# The source data. In production this is S3 or HDFS; here it is a local
# directory so the example runs. The Spark code that reads it is identical.
spark = SparkSession.builder.appName("SeedData").master("local[2]").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")
spark.createDataFrame([
    ("Sharding", "scaling",
     "Sharding splits one logical table across many physical machines."),
    ("Replication", "availability",
     "Replication keeps copies of the same data on more than one node."),
    ("Partitioning", "scaling",
     "Partitioning divides a table into smaller pieces inside one database."),
    ("Write-ahead log", "durability",
     "A write-ahead log records changes before they are applied to pages."),
], ["title", "category", "content"]).write.mode("overwrite").parquet("documents.parquet")
spark.stop()
print("wrote documents.parquet")

client = MilvusClient(uri="http://localhost:19530")
if client.has_collection("spark_documents"):
    client.drop_collection("spark_documents")

schema = MilvusClient.create_schema(auto_id=True, enable_dynamic_field=False)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=384)
schema.add_field("title", DataType.VARCHAR, max_length=1000)
schema.add_field("category", DataType.VARCHAR, max_length=50)

index_params = client.prepare_index_params()
index_params.add_index(field_name="embedding", index_type="HNSW",
                       metric_type="COSINE", params={"M": 16, "efConstruction": 256})

client.create_collection(collection_name="spark_documents", schema=schema,
                         index_params=index_params)
print("created collection spark_documents")

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, FloatType

spark = SparkSession.builder.appName("VectorIndexing").master("local[2]").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")

df = spark.read.parquet("documents.parquet")

# Cache the model per Python worker. Loading it inside the UDF body without
# this global would reload it for every single row.
_model = None


def get_model():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer

        _model = SentenceTransformer("all-MiniLM-L6-v2")
    return _model


@udf(returnType=ArrayType(FloatType()))
def generate_embedding(text):
    return get_model().encode(text).tolist()


df_with_embeddings = df.withColumn("embedding", generate_embedding(df.content))


# One client per partition, not one per row: connections are the thing that
# does not survive being created a million times.
def upsert_to_milvus(partition_iter):
    from pymilvus import MilvusClient

    rows = [
        {"embedding": row.embedding, "title": row.title, "category": row.category}
        for row in partition_iter
    ]
    if rows:
        milvus = MilvusClient(uri="http://localhost:19530")
        milvus.insert(collection_name="spark_documents", data=rows)
        milvus.close()


df_with_embeddings.foreachPartition(upsert_to_milvus)
spark.stop()

from pymilvus import MilvusClient

client = MilvusClient(uri="http://localhost:19530")
client.flush("spark_documents")
print("Batch indexing complete")
print(client.get_collection_stats("spark_documents"))