# Rate limiter decorator
@sleep_and_retry
@limits(calls=1500, period=60)
def get_embedding(text):
    try:
        result = genai.embed_content(
            model="models/text-embedding-004",
            content=text,
            task_type="retrieval_document",
        )
        return result['embedding']
    except Exception as e:
        print(f"Error getting embedding: {e}")
        return None

def add_embeddings_to_products(batch_size=50):
    # Get the total number of products to process
    total_query = """
    MATCH (p:Product)
    WHERE p.description_embedding IS NULL
      AND p.description IS NOT NULL
    RETURN count(p) AS total
    """
    total_result = graph.run(total_query).data()
    total_to_process = total_result[0]['total'] if total_result else 0
    print(f"Total products to process: {total_to_process}n")
    total_processed = 0
    # Initialize tqdm progress bar
    with tqdm(total=total_to_process, desc='Processing products', unit='product') as pbar:
        while True:
            # Get batch of products
            query = """
            MATCH (p:Product)
            WHERE p.description_embedding IS NULL
              AND p.description IS NOT NULL
            RETURN p.uniq_id AS id, p.description AS description
            LIMIT $batch_size
            """
            products = graph.run(query, parameters={'batch_size': batch_size}).data()
            if not products:
                break
            # Process each product in the batch
            for product in products:
                try:
                    if product['description']:
                        embedding = get_embedding(product['description'])
                        if embedding:
                            # Update product with embedding
                            graph.run("""
                            MATCH (p:Product {uniq_id: $id})
                            SET p.description_embedding = $embedding
                            """, parameters={
                                'id': product['id'],
                                'embedding': embedding
                            })
                    total_processed += 1
                    pbar.update(1)  # Update the progress bar
                except Exception as e:
                    print(f"Error processing product {product['id']}: {e}")
            # Add a small delay between batches
            time.sleep(1)
    print(f"nTotal products processed: {total_processed}")
    return total_processed
# Add embeddings to products
print("Adding embeddings to products...n")
total_processed = add_embeddings_to_products()
print(f"nProcess completed. Total products processed: {total_processed}")