tools/genai-prod-catalog-enrichment/main.py

import os
import json
import csv
import traceback
import datetime
import time
import requests
import fitz
from pdf_helper import parse_pdf
from genai_helper import generate_faqs, \
    generate_isqs, generate_company_details, \
    get_specific_caption, \
    generate_tags_json, generate_category_json
from genai_helper import generate_images_json as gij
from text_helper import clean_text, get_company_text
from google.cloud import storage


def get_todays_date():
    """Returns today's date in 'DD_MM_YYYY' format."""
    today = datetime.date.today()
    return today.strftime('%d_%m_%Y')


def download_blob(bucket_name, source_blob_name,
                  destination_file_name, prefix, csv_folder):
    """Downloads a blob from the bucket."""

    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(f"{prefix}/{source_blob_name}")
    blob.download_to_filename(f"{csv_folder}/{destination_file_name}")
    print(
        "[INFO]: Downloaded storage object "
        "{} from bucket {} to local file {}.".format(
            source_blob_name, bucket_name, destination_file_name
        )
    )


def upload_blob(bucket_name, source_file_name,
                destination_blob_name, prefix, date_):
    """Uploads a file to the bucket."""
    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(f"{prefix}/{date_}/{destination_blob_name}")
    blob.upload_from_filename(source_file_name)

    print(f"[INFO]: File {source_file_name} "
          f"uploaded to {destination_blob_name}.")


def get_csv_info(csv_gcs_uri, bucket_name, csv_folder):
    download_blob(bucket_name, csv_gcs_uri.split('/')[-1],
                  csv_gcs_uri.split("/")[-1],
                  csv_gcs_uri.split('/')[-2], csv_folder)
    with open(f"{csv_folder}/{csv_gcs_uri.split('/')[-1]}", "r") as f:
        reader = csv.reader(f)
        csv_data = list(reader)
    return csv_data[1:]


def download_pdf(url, filename):
    response = requests.get(url)
    with open(filename, 'wb') as f:
        f.write(response.content)


def update_output_json(json_output_response, faq_json, isq_json,
                       product_images, product_category, product_tags, error_ws2_message):
    todays_date = get_todays_date()
    products_output_response = get_output_template()
    products_output_response[0]["product_name_citation"][0]["created_date"] = todays_date
    products_output_response[0]["specification_citation"][0]["created_date"] = todays_date
    products_output_response[0]["product_image_citation"][0]["created_date"] = todays_date
    products_output_response[0]["product_category_citation"][0]["created_date"] = todays_date

    try:
        products_output_response[0]["product_images"] = product_images
        products_output_response[0]["product_category"] = product_category
        products_output_response[0]["product_tags"] = product_tags
        if "response_error" not in isq_json:
            if "product_name" in isq_json:
                products_output_response[0]["product_name"] = isq_json["product_name"]
            if "confidence_score" in isq_json:
                products_output_response[0]["product_name_citation"][0]["confidence_score"] = isq_json["confidence_score"]
                products_output_response[0]["specification_citation"][0]["confidence_score"] = isq_json["confidence_score"]
            if "specifications" in isq_json:
                products_output_response[0]["specification"] = isq_json["specifications"]
            json_output_response["products"].extend(products_output_response)
        else:
            products_output_response[0] = {**products_output_response[0], **isq_json}
            json_output_response["products"].extend(products_output_response)
        if "response_error" not in products_output_response[0] and error_ws2_message != '':
            response_error = {"message": error_ws2_message}
            products_output_response[0]["response_error"] = response_error
        if "response_error" not in faq_json:
            json_output_response["catalogue_faqs"].extend(faq_json[list(faq_json.keys())[0]])
        else:
            json_output_response["catalogue_faqs"].append(faq_json)

        return json_output_response

    except Exception as e:
        print(f"[ERROR]: Error while updating output JSON - {e}")
        return {}


def get_output_template():
    products_output_response = [
        {
            "product_name": [],
            "product_name_citation": [
                {
                    "pdf_name": "",
                    "url": "",
                    "page": 0,
                    "startIndex": [],
                    "endIndex": [],
                    "confidence_score": 0.0,
                    "created_date": "01/01/1970"
                }
            ],
            "specification": {},
            "specification_citation": [
                {
                    "attribute": "",
                    "pdf_name": "",
                    "url": "",
                    "page": 0,
                    "startIndex": [],
                    "endIndex": [],
                    "confidence_score": 0.0,
                    "created_date": "01/01/1970"
                }
            ],
            "product_images": [
            ],
            "product_image_citation": [
                {
                    "pdf_name": "",
                    "url": "",
                    "page": 0,
                    "startIndex": [],
                    "endIndex": [],
                    "confidence_score": 0.0,
                    "created_date": "01/01/1970"
                }
            ],
            "product_category": [],
            "product_tags": [],
            "product_category_citation": [
                {
                    "pdf_name": "",
                    "url": "",
                    "page": 0,
                    "startIndex": [],
                    "endIndex": [],
                    "confidence_score": 0.0,
                    "created_date": "01/01/1970"
                }
            ]
        }
    ]
    return products_output_response


def create_json_from_dict(json_path, dict_, gcs_bucket=None):
    if gcs_bucket:
        bucket_object = storage.Client().get_bucket(gcs_bucket)
        json_path = bucket_object.blob(json_path)
        json_obj = json.dumps(dict_)
        json_path.upload_from_string(data=json_obj,
                                     content_type='application/json')
    else:
        with open(json_path, "w") as outfile:
            json.dump(dict_, outfile)


def check_pdf_type(filename, input_gcs_bucket=None):
    scanned = False
    tables = False
    images = False
    total_images = 0
    if input_gcs_bucket:
        bucket_object = storage.Client().bucket(input_gcs_bucket)
        blob = bucket_object.blob(filename)
        pdf_file = fitz.open("pdf", blob.download_as_bytes())
    else:
        pdf_file = fitz.open(filename)
    pages = []
    for page_no, page in enumerate(pdf_file, start=1):
        page_wise_details = {"scanned": False,
                             "Images": False,
                             "Tables": False,
                             "Total Images": 0}
        print(f"Page no {page_no}")
        if not page.get_text():
            print("[Info]: scanned!")
            scanned = True
            page_wise_details["scanned"] = True
        images_info = page.get_images()
        print(f"[Info]: No of images: {len(images_info)}")
        if len(images_info) > 0:
            images = True
            page_wise_details["Images"] = True
            page_wise_details["Total Images"] = len(images_info)
            total_images = total_images + len(images_info)
        tables = page.find_tables()
        if len(list(tables)) > 0:
            print("[Info]: Tables present")
            tables = True
            page_wise_details["Tables"] = True
        pages.append(page_wise_details)
    return {"scanned": scanned, "Tables": tables,
            "Images": images, "Total_Images": total_images,
            "pages": pages}


def end_to_end_pipeline(input_pdf_uri, output_gcs_bucket, project_id):
    # csv_folder = "csv"
    pdf_folder = "pdf_files"
    # date_ = get_todays_date()
    # os.makedirs(csv_folder, exist_ok=True)
    os.makedirs(pdf_folder, exist_ok=True)

    try:
        start = time.time()
        # csv_data = get_csv_info(csv_gcs_uri, bucket_name, csv_folder)
        end = time.time()
        print(f"[INFO]: CSV data fetched "
              f"successfully in {end - start} seconds")
    except Exception as e:
        print(f"[ERROR]: Error "
              f"during fetching data from CSV - {e}")
        return None

    uri = input_pdf_uri

    # check if file name exists
    full_path = uri.replace("gs://", "")
    input_gcs_bucket = full_path.split("/")[0]
    filename = full_path.replace(f"{input_gcs_bucket}/", "")
    name = filename.split("/")[-1]
    output_gcs_bucket = output_gcs_bucket

    try:
        pdf_type = check_pdf_type(filename, input_gcs_bucket)
        if not pdf_type["scanned"]:
            start = time.time()
            pdf_json = parse_pdf(filename, input_gcs_bucket, output_gcs_bucket)
            end = time.time()
            print(f"[INFO]: Parsed PDF successfully in {end - start} seconds")

            pdf_json = get_specific_caption(pdf_json)
            name_initials = name.replace(".pdf", "")
            json_path = f"./{name_initials}.json"
            create_json_from_dict(json_path, pdf_json, output_gcs_bucket)
            pdf_json = clean_text(pdf_json)
            json_output_response = {
                "pdf_name": f"{filename.split('/')[-1]}",
                "pdf_url": uri,  # f"gs://{input_gcs_bucket}/{filename}",
                "company_details": {},
                "products": [],
                "catalogue_faqs": []
                # "pc_item_doc_id": pc_item_doc_id,
                # "pc_doc_modified_date": pc_doc_modified_date,
                # "pc_item_doc_path": pc_item_doc_path,
                # "fk_pc_item_id": fk_pc_item_id,
                # "fk_glusr_usr_id": fk_glusr_usr_id
            }
            print(json_output_response, "\n")
            start = time.time()
            company_text = get_company_text(pdf_json['pages'])
            if company_text:
                company_details = \
                    generate_company_details(company_text, project_id)
                if "company_details" in company_details:
                    json_output_response["company_details"] = \
                        company_details["company_details"]
                elif "response_error" in company_details:
                    json_output_response["company_details"] = \
                        company_details
                else:
                    json_output_response["company_details"] = {}
            end = time.time()
            print(f"[INFO]: Company details "
                  f"extraction time: {end - start} seconds")
            for page_no, page in enumerate(pdf_json['pages']):
                try:
                    print(str(page_no))
                    context = page['texts']['full_text']
                    start = time.time()
                    faq_json = generate_faqs(context, project_id)
                    isq_json = generate_isqs(context, project_id)
                    end = time.time()
                    print(f"[INFO]: FAQ, ISQ generation "
                          f"time for page {page_no + 1}: "
                          f"{end - start} seconds")

                    error_ws2_msg = ''
                    try:
                        products = isq_json['product_name']
                    except Exception:
                        products = []
                    print("Products:", products)
                    if len(products) > 0:
                        try:
                            start = time.time()
                            product_tags = \
                                generate_tags_json(context, products)
                            end = time.time()
                            print(f"[INFO]: Product tag "
                                  f"generation time for page {page_no + 1}: "
                                  f"{end - start} seconds")
                        except Exception:
                            print(f"[ERROR]: Error during generation "
                                  f"of product tags -"
                                  f" {str(traceback.format_exc())}")
                            error_ws2_msg = error_ws2_msg
                            product_tags = {}
                        try:
                            start = time.time()
                            product_category = \
                                generate_category_json(context, products)
                            end = time.time()
                            print(
                                f"[INFO]: Product category generation time "
                                f"for page {page_no + 1}: "
                                f"{end - start} seconds")
                        except Exception:
                            print(
                                f"[ERROR]: Error during generation of"
                                f" product category - "
                                f"{str(traceback.format_exc())}")
                            product_category = {}

                        try:
                            start = time.time()
                            image_json = \
                                gij(page, products,
                                    product_tags, output_gcs_bucket)
                            end = time.time()
                            print(f"[INFO]: Image JSON generation "
                                  f"time for page {page_no + 1}:"
                                  f" {end - start} seconds")
                        except Exception as error:
                            print(
                                f"[ERROR]: Error during generation "
                                f"of product category - "
                                f"{str(traceback.format_exc())} + {error}")
                            error_ws2_msg = \
                                error_ws2_msg + "\n" \
                                + 'Product image json ' \
                                  'generation failed - {str(e)}'
                            image_json = {}
                    else:
                        product_tags = {}
                        product_category = {}
                        image_json = {}

                    start = time.time()
                    json_output_response = \
                        update_output_json(json_output_response,
                                           faq_json, isq_json, image_json,
                                           product_category,
                                           product_tags, error_ws2_msg)
                    end = time.time()
                    print(f"[INFO]: Output for page {page_no + 1}: "
                          f"updated in {end - start} seconds")
                except Exception as error:
                    print(f"[ERROR]: Error during generating "
                          f"final json - "
                          f"{str(traceback.format_exc())} - {error}")

            final_json_path = f"{name_initials}_ouput.json"
            start = time.time()
            create_json_from_dict(final_json_path,
                                  json_output_response,
                                  output_gcs_bucket)
            end = time.time()
            print(f"[INFO]: Final output JSON "
                  f"uploaded to GCS in {end - start} seconds")
        else:
            print("[INFO]: Please process text PDF's only!")

        return json_output_response
    except Exception as error:
        print(f"[ERROR]: Error for filename: "
              f"{filename} - {str(traceback.format_exc())} - {error}")