demo.py

"""
demo.py — Walkthrough of VectorDB: embed → store → search.

Run:  python demo.py
"""

from vector_db import VectorDB


# -- 1. Sample documents -----------------------------------------------
#
# A small, diverse corpus so we can see semantic search working
# (not just keyword overlap).

DOCUMENTS = [
    # Tech
    ("Python is a high-level programming language known for its readability and simplicity.",
     {"category": "tech", "topic": "programming"}),
    ("Machine learning models learn patterns from data without being explicitly programmed.",
     {"category": "tech", "topic": "AI/ML"}),
    ("Docker containers package an application and its dependencies into a single unit.",
     {"category": "tech", "topic": "devops"}),
    ("A REST API allows different software systems to communicate over HTTP.",
     {"category": "tech", "topic": "backend"}),
    ("Vector databases store high-dimensional embeddings and support similarity search.",
     {"category": "tech", "topic": "database"}),

    # Science
    ("Black holes are regions of spacetime where gravity is so strong that nothing can escape.",
     {"category": "science", "topic": "astronomy"}),
    ("DNA carries the genetic instructions for the development and functioning of all living things.",
     {"category": "science", "topic": "biology"}),
    ("Photosynthesis is the process by which plants convert sunlight into chemical energy.",
     {"category": "science", "topic": "biology"}),
    ("Quantum entanglement links two particles so that the state of one instantly affects the other.",
     {"category": "science", "topic": "physics"}),

    # Food
    ("Sushi is a Japanese dish made of vinegared rice paired with raw fish or other toppings.",
     {"category": "food", "topic": "Japanese cuisine"}),
    ("Pizza originated in Naples and consists of a dough base topped with tomato sauce and cheese.",
     {"category": "food", "topic": "Italian cuisine"}),
    ("A smoothie is a blended drink made from fruits, vegetables, and often yogurt or milk.",
     {"category": "food", "topic": "drinks"}),

    # History
    ("The Roman Empire was one of the largest empires in ancient history, spanning three continents.",
     {"category": "history", "topic": "Rome"}),
    ("The Renaissance was a cultural movement in Europe that revived interest in art and classical learning.",
     {"category": "history", "topic": "Renaissance"}),

    # Sports
    ("Football (soccer) is the world's most popular sport, played by over 250 million people.",
     {"category": "sports", "topic": "soccer"}),
    ("The Olympic Games are held every four years and bring athletes from all over the world together.",
     {"category": "sports", "topic": "Olympics"}),
]


# -- 2. Build and save the database -----------------------------------

def build_db():
    db = VectorDB(db_path="data.json")
    db._records = []   # always rebuild from scratch so re-runs don't duplicate
    print(f"Adding {len(DOCUMENTS)} documents and embedding each one...\n")
    for text, meta in DOCUMENTS:
        db.add(text, meta)
    db.save()
    return db


# -- 3. Search ---------------------------------------------------------

def run_searches(db: VectorDB):
    queries = [
        "how do computers learn from data?",
        "space and the universe",
        "food from Italy",
        "athletic competition between nations",
        "storing and querying high-dimensional data",
    ]

    separator = "-" * 64

    for query in queries:
        print(f"\n{separator}")
        print(f"  Query: \"{query}\"")
        print(separator)

        results = db.search(query, top_k=3)
        for rank, r in enumerate(results, start=1):
            score_bar = "#" * int(r["score"] * 20)        # visual bar
            print(f"  #{rank}  score={r['score']:.4f}  {score_bar}")
            print(f"       [{r['metadata'].get('category','?')} / {r['metadata'].get('topic','?')}]")
            print(f"       {r['text']}")

    print(f"\n{separator}\n")


# -- 4. Main -----------------------------------------------------------

if __name__ == "__main__":
    print("=" * 64)
    print("  Vector DB Demo")
    print("  all-MiniLM-L6-v2 embeddings  |  cosine similarity  |  JSON store")
    print("=" * 64 + "\n")

    # Build from scratch the first time; or reload if data.json already exists.
    db = build_db()

    print("\nNow searching across the stored vectors...")
    run_searches(db)

    print("Tip: open data.json to see how the 384-dimensional")
    print("     embedding vectors are stored alongside each document.\n")