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:
    def __init__(self, stars: int, review_text: str):
        self.stars = stars
        self.review_text = review_text


class Product:
    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
    """
    @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):
    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:
    output_file_name = f"products_{category.lower().replace(' ', '_')}.json"
    return os.path.join(data_dir, output_file_name)


def category_review_file(category: str) -> str:
    output_file_name = f"reviews_{category.lower().replace(' ', '_')}.json"
    return os.path.join(data_dir, output_file_name)


def products_for_category(category: str) -> List[Product]:
    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]:
    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:
    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]]:
    product_features_file = os.path.join(data_dir, '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):
    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():
    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):
    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):
    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):
    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):
    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]))