src/data_synthesis/generate_data.py

"""
This script contains the first step of data synthesis - generation of the json files containing
all the product categories, product features and synthetic product reviews.  The reuslt of this script
is the generation of the json files in the data directory.
"""

import json
import openai
import os
import sys
import time
from typing import Dict, List

from src.common import data_dir


class Review:
    """
    Simple representation of a user Review of a Product
    """
    def __init__(self, stars: int, review_text: str):
        self.stars = stars
        self.review_text = review_text


class Product:
    """
    Simple representation of a prduct
    """
    def __init__(self, category: str, name: str, description: str, price: float, features: List[str], reviews: List[Review]):
        self.category = category
        self.name = name
        self.description = description
        self.price = price
        self.features = features
        self.reviews = reviews


class DataPrompt:
    """
    Holder for static prompt generation functions for the data generation process
    """
    @staticmethod
    def prompt_setup() -> str:
        return "You are a marketing assistant for consumer home electronics manufacturer ElectroHome. You are polite and succinct.\n\n"


    @staticmethod
    def prompt_setup_user() -> str:
        return "You are a customer of consumer home electronics manufacturer ElectroHome, and are reviewing a product you have purchased and used.\n\n"


    @staticmethod
    def products_for_category(category: str, features: List[str], k: int) -> str:
        existing_products = product_names_for_category(category)
        prompt = f"Suggest exactly {k} products in the category {category}. \nPlease give the products realistic product names but cover a range of different customer needs (e.g budget, premium, compact, eco, family).\nDo not include the customer needs words in the product name.\nProduct names must be unique."
        if len(existing_products) > 0:
            prompt += f" The following product names are already in use, so do not duplicate them: {', '.join(existing_products)}"
        prompt += "\nPlease select between 4 and 8 features for each product from the following options: {', '.join(features)}.\n"
        prompt += """
Please format the response as json in this style:
{
  "products": [
    {
      "name": "product name",
      "features": ["feature 1", "feature 2"],
      "price": "$49.99",
      "description": "A description of the product in 50 to 100 words."
    }
  ]
}"""
        return prompt

    @staticmethod
    def format_features(features: List[str]) -> str:
        """
        Convenience method to do comma/and join
        """
        if len(features) == 0:
            return ""
        if len(features) == 1:
            return features[0]
        return (', '.join(features[:-1])) + f' and {features[-1]}'


    @staticmethod
    def reviews_for_product(product: Product, k: int):
        prompt = f"Suggest exactly {k} reviews for this product.\nThe product is a {product.category.lower()[0:-1]} named the '{product.name}', which features {DataPrompt.format_features(product.features)}.\nFirst pick an integer star rating from 1 to 5 stars, where 1 is bad and 5 is great, for the review.\nNext write the review text of between 50 and 100 words for the review from the user. The text in the review should align to the star rating, so if the rating is 1 the review would be critical and if the rating is 5 the review would be positive.\n"
        prompt += """
Please format the response as json in this style:
{
  "reviews": [
    {
      "stars": 3,
      "review_text": "Between 50 and 100 words reviewing the product go here."
    }
  ]
}"""
        return prompt


def generate_products(category: str, features: List[str], k: int = 20):
    """
    Invoke GPT3.5 Turbo model and get it to generate some products based on a category
    """
    prompt = DataPrompt.products_for_category(category, features, k)
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-16k",
        messages=[
            {"role": "system", "content": DataPrompt.prompt_setup()},
            {"role": "user", "content": prompt}
        ],
        temperature=1.0
    )
    output_text = response['choices'][0]['message']['content']
    add_products(category, output_text, k)


def category_product_file(category: str) -> str:
    """
    Utility to get the file containing products in a category
    """
    output_file_name = f"products_{category.lower().replace(' ', '_')}.json"
    return os.path.join(data_dir, 'json', output_file_name)


def category_review_file(category: str) -> str:
    """
    Utility to get the file containing reviews of products in a category
    """
    output_file_name = f"reviews_{category.lower().replace(' ', '_')}.json"
    return os.path.join(data_dir, 'json', output_file_name)


def products_for_category(category: str) -> List[Product]:
    """
    Load all the associated products which have been generated for this
    category, and the reviews, then merge the two and return a list of
    all the products in this category along with their reviews
    """
    cat_file = category_product_file(category)
    if not os.path.exists(cat_file):
        return []
    else:
        products = []
        with open(cat_file, 'r') as f:
            category_json = json.load(f)
            for prod in category_json['products']:
                price = float(prod['price'][1:])
                p = Product(category, prod['name'], prod['description'], price, prod['features'], [])
                products.append(p)
        reviews_file = category_review_file(category)
        if os.path.exists(reviews_file):
            with open(reviews_file, 'r') as f:
                review_json = json.load(f)
                for p in products:
                    if p.name in review_json:
                        for review in review_json[p.name]:
                            p.reviews.append(Review(review['stars'], review['review_text']))
        return products


def product_names_for_category(category: str) -> List[str]:
    """
    Get a list of just the names of the products in this category
    from the generated product json file
    """
    cat_file = category_product_file(category)
    if not os.path.exists(cat_file):
        return []
    else:
        names = []
        with open(cat_file, 'r') as f:
            category_json = json.load(f)
            for prod in category_json['products']:
                names.append(prod['name'])
            return names


def add_products(category: str, product_json: str, k: int) -> None:
    """
    Given a string of json representing newly generated products,
    add those products to the existing product json file for this category
    """
    cat_file = category_product_file(category)
    if not os.path.exists(cat_file):
        with open(cat_file, 'w') as f:
            f.write(product_json)
    else:
        with open(cat_file, 'r') as f:
            existing_products = json.load(f)
            new_products = json.loads(product_json)
            count = 0
            for new_p in new_products['products']:
                if count >= k:
                    break
                existing_products['products'].append(new_p)
                count += 1
        with open(cat_file, 'w') as f:
            json.dump(existing_products, f, indent=2)


def get_categories_and_features() -> Dict[str, List[str]]:
    """
    Get dictionary of will each category as a key and the list of available
    features to products in that category as the value
    """
    product_features_file = os.path.join(data_dir, 'json', 'product_features.json')
    cats_and_feats = {}
    with open(product_features_file, 'r') as f:
        feature_json = json.load(f)
        for cat in feature_json['categories']:
            cat_name = cat['category']
            cat_features = cat['features']
            cats_and_feats[cat_name] = cat_features
    return cats_and_feats


def generate_all_products(target_count=40):
    """
    Generate all products for all categories, trying to reach a given target count
    of products.
    """
    product_features_file = os.path.join(data_dir, 'product_features.json')

    with open(product_features_file, 'r') as f:
        feature_json = json.load(f)
        for cat in feature_json['categories']:
            cat_name = cat['category']
            cat_features = cat['features']
            existing_products = product_names_for_category(cat_name)
            if len(existing_products) < target_count:
                num_to_generate = target_count - len(existing_products)
                print(f"Generating {num_to_generate} {cat_name}")
                generate_products(cat_name, cat_features, num_to_generate)
            else:
                print(f"Skipping {cat_name} as targetting {target_count} and already have {len(existing_products)}")


def dump_products_to_csv():
    """
    Dump a csv file for debug, for every product showing category name and product name
    """
    cats = get_categories_and_features().keys()
    cat_keys = []
    for cat in cats:
        for prod in product_names_for_category(cat):
            cat_keys.append(f"{cat},{prod}")
    dump_file = os.path.join(data_dir, "products.csv")
    with open(dump_file, 'w') as f:
        f.write('\n'.join(cat_keys))


def generate_reviews(target_count: int):
    """
    Generate reviews for each category up to a target count of reviews
    """
    for cat in get_categories_and_features().keys():
        generate_reviews_for_category(cat, target_count)


def generate_reviews_for_category(category: str, target_count: int):
    """
    Generate reviews for a specific category up to a given target number of reviews
    """
    batch_size = 25  # Max number of reviews to request in one go from GPT

    # Set up a loop to continue trying to find more work to do until complete
    working = True
    recent_exception = False
    while working:
        working = False
        products = products_for_category(category)
        for prod in products:
            if len(prod.reviews) < target_count:
                working = True
                reviews_to_request = min([batch_size, target_count - len(prod.reviews)])
                try:
                    print(f'{prod.category[:-1]}: {prod.name} has {len(prod.reviews)} reviews. Requesting {reviews_to_request} more.')
                    generate_reviews_for_product(prod, reviews_to_request)
                    recent_exception = False
                except openai.error.ServiceUnavailableError:
                    print("GPT is overloaded - waiting 10 seconds.....")
                    recent_exception = False
                    time.sleep(10)
                except Exception as e:
                    print(f"Exception {e} in generating reviews")
                    if recent_exception:
                        print(f"Exception appears to be stubborn so throwing out")
                        raise e
                    recent_exception = True
            else:
                print(f'{prod.category[:-1]}: {prod.name} has {len(prod.reviews)} reviews ({target_count} requested). Skipping.')


def generate_reviews_for_product(product: Product, k: int):
    """
    Generate a number of reviews from GPT3.5 for a specific product and add them to the product
    """
    prompt = DataPrompt.reviews_for_product(product, k)
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-16k",
        messages=[
            {"role": "system", "content": DataPrompt.prompt_setup_user()},
            {"role": "user", "content": prompt}
        ],
        temperature=1.0
    )
    output_text = response['choices'][0]['message']['content']
    add_reviews_to_product(output_text, product)


def add_reviews_to_product(reviews_json: str, product: Product):
    """
    Load the reviews file containing this product category, append this review to the list and
    re-save the file
    """
    reviews_json = json.loads(reviews_json)
    reviews_file = category_review_file(product.category)
    if not os.path.exists(reviews_file):
        category_data = {product.name: reviews_json['reviews']}
        with open(reviews_file, 'w') as f:
            json.dump(category_data, f, indent=2)
    else:
        with open(reviews_file, 'r') as f:
            existing_reviews = json.load(f)
            if product.name in existing_reviews:
                for r in reviews_json['reviews']:
                    existing_reviews[product.name].append(r)
            else:
                existing_reviews[product.name] = reviews_json['reviews']
        with open(reviews_file, 'w') as f:
            json.dump(existing_reviews, f, indent=2)


"""
# The sequence of steps to arrive at the final JSON files containing the data is as follows:
# Manual step - generated product categories and product features from GPT and loaded to file
# run generate_all_products()  # Generate 40 products in each category
# run dump_products_to_csv()  # Dump the products to csv for manual name check
# Manual step - review names and tweak some of them directly in the json files
# run generate_reviews_for_category(50) for each category  # Generate 50 reviews per product in every category
"""
if __name__ == "__main__":
    generate_reviews_for_category(sys.argv[1], int(sys.argv[2]))