import numpy as np

def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Angle between vectors. Range: -1 to 1. 1 = identical direction."""
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

def dot_product(a: list[float], b: list[float]) -> float:
    """Dot product. Equivalent to cosine if vectors are normalised."""
    return float(np.dot(np.array(a), np.array(b)))

def euclidean_distance(a: list[float], b: list[float]) -> float:
    """Straight-line distance. Lower = more similar. Range: 0 to inf."""
    return float(np.linalg.norm(np.array(a) - np.array(b)))

# Demonstrate: semantically similar texts should be close
vecs = embed_openai([
    "DPDK is a fast packet processing framework",
    "FD.io DPDK provides high-speed networking",
    "Machine learning uses gradient descent optimisation"
])

print(cosine_similarity(vecs[0], vecs[1]))  # ~0.91 — very similar
print(cosine_similarity(vecs[0], vecs[2]))  # ~0.18 — very different

pip install chromadb openai

import chromadb
from chromadb.utils import embedding_functions

# ── In-memory (for tests / notebooks) ────────────────
client = chromadb.Client()

# ── Persistent (survives restarts) ───────────────────
client = chromadb.PersistentClient(path="./chroma_db")

# ── Use OpenAI embeddings automatically ──────────────
oai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.environ["OPENAI_API_KEY"],
    model_name="text-embedding-3-small"
)

# Create or get a collection
collection = client.get_or_create_collection(
    name="docs",
    embedding_function=oai_ef,           # auto-embeds on add/query
    metadata={"hnsw:space": "cosine"}   # use cosine similarity
)

# Add documents — Chroma embeds them automatically
collection.add(
    ids=["doc1", "doc2", "doc3"],
    documents=[
        "DPDK is a set of libraries for fast packet processing",
        "VPP uses DPDK for high-performance networking in telecom",
        "Python is a general-purpose programming language"
    ],
    metadatas=[
        {"source": "dpdk_docs", "year": 2024},
        {"source": "vpp_docs",  "year": 2024},
        {"source": "python_docs", "year": 2023},
    ]
)

# Query — semantic search
results = collection.query(
    query_texts=["how does packet processing work?"],
    n_results=2,
    include=["documents", "metadatas", "distances"]
)
for doc, meta, dist in zip(
    results["documents"][0],
    results["metadatas"][0],
    results["distances"][0]
):
    print(f"Score: {1-dist:.3f} | Source: {meta['source']} | {doc[:60]}")

pip install pinecone-client

from pinecone import Pinecone, ServerlessSpec

pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])

# Create index (one-time setup)
pc.create_index(
    name="my-rag-index",
    dimension=1536,            # must match embedding model dimension
    metric="cosine",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)

index = pc.Index("my-rag-index")

# Upsert vectors (create or update)
vectors = embed_batch(documents)
index.upsert(vectors=[
    {
        "id":     f"doc_{i}",
        "values": vec,
        "metadata": {"text": doc, "source": "docs", "chunk_idx": i}
    }
    for i, (vec, doc) in enumerate(zip(vectors, documents))
])

# Query
query_vec = embed_openai(["packet processing performance"])[0]
results = index.query(
    vector=query_vec,
    top_k=5,
    include_metadata=True,
    filter={"source": {"$eq": "docs"}}
)

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

# Index stats
print(index.describe_index_stats())

# Install pgvector extension in PostgreSQL
# docker run -e POSTGRES_PASSWORD=pass -p 5432:5432 pgvector/pgvector:pg16

pip install psycopg2-binary pgvector

import psycopg2
from pgvector.psycopg2 import register_vector
import numpy as np

conn = psycopg2.connect("postgresql://postgres:pass@localhost/ragdb")
register_vector(conn)

# Enable extension and create table
with conn.cursor() as cur:
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
    cur.execute("""
        CREATE TABLE IF NOT EXISTS documents (
            id      SERIAL PRIMARY KEY,
            content TEXT NOT NULL,
            source  TEXT,
            embedding vector(1536)
        )
    """)
    cur.execute("CREATE INDEX IF NOT EXISTS docs_embedding_idx ON documents USING ivfflat (embedding vector_cosine_ops)")
    conn.commit()

# Insert documents with embeddings
def insert_docs(texts: list[str], source: str):
    vecs = embed_batch(texts)
    with conn.cursor() as cur:
        cur.executemany("""
            INSERT INTO documents (content, source, embedding)
            VALUES (%s, %s, %s)
        """, [(text, source, vec) for text, vec in zip(texts, vecs)])
    conn.commit()

# Semantic search — pure SQL!
def semantic_search(query: str, top_k: int = 5, source: str = None) -> list[dict]:
    query_vec = embed_openai([query])[0]
    source_filter = "AND source = %s" if source else ""
    params = [query_vec, top_k] if not source else [query_vec, source, top_k]

    with conn.cursor() as cur:
        cur.execute(f"""
            SELECT content, source,
                   1 - (embedding <=> %s::vector) AS similarity
            FROM documents
            {f"WHERE source = %s" if source else ""}
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, [query_vec] + ([source] if source else []) + [query_vec, top_k])
        rows = cur.fetchall()
    return [{"content": r[0], "source": r[1], "similarity": r[2]} for r in rows]

# pgvector distance operators:
# <->   Euclidean distance
# <=>   Cosine distance (1 - cosine_similarity)
# <#>   Negative dot product