AI Search on YugabyteDB with Native pgvector Integration + GCP

				
					import os
import random
import psycopg2
from google import genai
from dotenv import load_dotenv

load_dotenv()

if not os.getenv("YB_HOST"):
    print("❌ ERROR: Could not find YB_HOST. Please check your .env file!")
    exit(1)

client = genai.Client(api_key=os.getenv("AI_STUDIO_KEY"))
YB_CONFIG = {
    "host": os.getenv("YB_HOST"),
    "port": os.getenv("YB_PORT", "5433"),
    "user": os.getenv("YB_USER", "yugabyte"),
    "password": os.getenv("YB_PASSWORD"),
    "database": os.getenv("YB_DB", "yugabyte")
}

categories = ["Electronics", "Outdoor", "Home Office", "Kitchen", "Fitness"]
adjectives = ["Professional", "Ultra-durable", "Eco-friendly", "Compact", "Luxury"]
products_base = [
    ("Noise-cancelling Headphones", "Perfect for focusing in loud offices or airplanes. Advanced active noise cancellation and spatial audio for a focused work environment."),
    ("Waterproof Tent", "Reliable shelter for heavy rain and stormy camping trips. A rugged, weatherproof shelter designed for high-altitude trekking and extreme cold."),
    ("Ergonomic Chair", "Supports your back during long hours of remote work. Adjustable lumbar support and breathable mesh designed for long-duration remote work."),
    ("Cast Iron Skillet", "Retains heat perfectly for gourmet cooking at home. A heavy-duty skillet that offers superior heat retention for professional searing."),
    ("Yoga Mat", "Extra thick cushioning for joint support during heavy exercise. Non-slip texture and high-density cushioning to protect joints during intense floor workouts.")
]

def generate_bulk_data(count=50):
    data = []
    for i in range(count):
        base_name, base_desc = random.choice(products_base)
        adj = random.choice(adjectives)
        cat = random.choice(categories)
        
        name = f"{adj} {base_name} Gen-{i}"
        price = round(random.uniform(15.0, 250.0), 2)
        full_desc = f"{base_desc} This {cat} item is {adj.lower()} and highly rated."
        
        data.append({"name": name, "desc": full_desc, "price": price, "cat": cat})
    return data

def seed_database():
    products = generate_bulk_data(50)
    conn = None 
    cur = None
    
    try:
        conn = psycopg2.connect(**YB_CONFIG, sslmode='disable')
        cur = conn.cursor()
        print(f"Connected. Starting ingestion of {len(products)} items...")

        for p in products:
            # Let the model output its native 3,072 dimensions!
            res = client.models.embed_content(
                model="gemini-embedding-001", 
                contents=p['desc']
            )
            vector = res.embeddings[0].values
            
            cur.execute("""
                INSERT INTO products (id, name, description, price, category, description_embedding)
                VALUES (gen_random_uuid(), %s, %s, %s, %s, %s)
            """, (p['name'], p['desc'], p['price'], p['cat'], vector))
            print(f"Inserted: {p['name']}")
        
        conn.commit()
        print("\n--- Seeding Complete! ---")
    except Exception as e:
        print(f"\n❌ Database Error: {e}")
    finally:
        if cur: cur.close()
        if conn: conn.close()

if __name__ == "__main__":
    seed_database()