app.py

import random
import traceback
import gradio as gr 
import numpy as np
import os
from langchain_core.output_parsers import JsonOutputParser

from langchain_openai.chat_models import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage, AIMessage
from langchain_anthropic import ChatAnthropic, ChatAnthropicMessages
from langchain_google_genai import ChatGoogleGenerativeAI
# from langchain_groq import ChatGroq
import openai
import google.generativeai as genai

from langchain import hub
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader, CSVLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.vectorstores import VectorStoreRetriever

import time 
import weave

from str2escaped import str2escaped


# feature_text = "Brand: Duckly. \nProduct name: Duck runner pro. \nKey properties: t-shirt, for running, sweat wicking, for marathon, 100% cotton."

feature_text = """Brand: Puma. 
Product name: Puma duffel bag.
Key features: light weight, only 1kg, black"""

garment_type = "all"

reference_text = """Baukasten-Anzughose Melwin, mint meliert"""

reference_text_0 = """NULLUS STUDIOS
Black Camellia Lapel Brooch Coat
Brushed wool-blend melton coat.

· Detachable brooch at notched lapel
· Button closure
· Welt pocket and flap pockets
· Padded shoulders
· Size 30 x 20 x 60 cm 
· Weight only 9kg
· Four-button surgeon's cuffs
· TextPatch at cuff
· Logo cutout at back collar
· Welt pockets at interior
· Full cupro satin lining
· Logo-engraved antiqued silver-tone hardware"""

structure_text = """[language = German]
{{ product_title | factual, short, one line of text, starting with 1 word defining the Style followed by the Product Name, then the Color separated by a comma }}
{% if the product fabric has a pattern - use the Images to determine the pattern in a single word after the Color %}
{% if the Style has no value - use the Images to determine the style of the product %}
""" 

structure_text_0 = \
"""# Headline {{ headline | inspiring, bold, action-oriented, max 8 words }}

## Introduction
{{ introduction_paragraph | motivational, passionate, 2-3 sentences }}

## Features and Benefits
{% for feature in features %}
### Feature {{ loop.index }}: {{ feature.name | dynamic, direct, 5-6 words }}
{{ feature.details | energetic, clear, 3-4 sentences }}
{% endfor %}

## Technical Specifications
{{ technical_specs | informative, to the point, concise list format }}
"""

structure_text_1 = """[type: UK website, style=true, language=English]
{{ introduction_paragraph | motivational, passionate, 1-2 sentences }}
{% for feature in features as bulleted list %}
{{ feature.description | dynamic, direct, 3-6 words }}
{% endfor %}
{{ technical_specs | informative, to the point, concise list format }}"""

structure_text_2 = """[type: Japanese newsletter, style=true, language=Japanese]
{{ introduction_paragraph | motivational, passionate, 3-6 sentences }}"""

guidance_prompt_text = """If it is made of 100% cotton, say "pure cotton" instead.
If it is lightweight, say "featherlight" instead."""

languages = ["American English",
             "British English",
             "German",
             "French",
             "Chinese",
             "Spanish",
             "Dutch",
             "Italian",
             "Japanese",
             "Polish",
             "Portuguese"]

models = ["gpt-4o",
          "claude-3-7-sonnet-latest",
          "claude-3-5-sonnet-20240620",
          "claude-3-5-sonnet-20241022",
          "gemini-2.0-flash-thinking-exp-01-21",
          ]

default_model = "gemini-2.0-flash-thinking-exp-01-21"

openai.api_key = os.environ["OPENAI_API_KEY"]

struct_copy_prompt = """Generate {nversions} versions of the product description for a product with the following information.
Write in a way that target the customer.
For each version, write at least {min_length} words and at most {max_length} words.
Make sure that the structure of each output follows the reference structure. Do not add any additional sentences or structures that are not in the reference structure.
Make sure to use the tone of voice, rythm, cadence and style of the reference copy for each output.
Use markdown format for each output.
Do not include any part of the reference structure in the output.
The structure of the output should follow the reference structure.
Do not use the structure of the reference copy in the output. 
Do not use any of the excluded words in the output.
Try to inlcude included words in the output when relevant.
Use the relevant information from the product features and intended use in the output.
Do not hallucinate any information about the product, use only the provided key features and intended use to write about the product.
Note that the reference copy should be used for style and tone only, do not use any part of the reference copy in the output. 
Do not hallucinate information about size and weight. Write about size and weight only if it is available in the list of features.
Use creative language in each output, do not use the common ways of starting product descriptions.
Avoid common phrases and cliches such as "Step into something", "Elevate something", "Discover something", "Unleash something", "Embrace something", and similar phrases.
For each version, try to write in different style.
Put each item in the list on a separate line. Use '-' character to start each item.  
Rate the quality of each version based on the following criteria:
    - how faithful it describes the product features. 
    - how well it follows the reference structure. 
    - how well it follows the tone of voice, rythm, cadence and style of the reference copy.
    - how well it avoid the excluded words.                                 
    - how well it includes the included words.                                 
    - how creative the language is.                                    
The score should be a number between 0 and 10 with 10 being the best quality.                              
Return the result in the following JSON format:
{{
    "versions": [                                 
        {{
            "id": 1,
            "content": The first product description,
            "explanation": A less than 20 word explanation of the score of the first product description,                         
            "score": The score of the first product description
        }},
        {{
            "id": 2,
            "content": The second product description,
            "explanation": A less than 20 word explanation of the score of the first product description,                                     
            "score": The score of the second product description
        }},
        ...                                 
    ],                                      
    "best_version": {{
        "explanation": Explanation for why this version is the best,                                                              
        "id": The id of the best version                                                             
    }}                                                          
}}                              
Make sure that the output is in JSON format, no extra text should be included in the output.                                
                                 
<product_information>
<key_features>
{key_features}
</key_features>

<reference_structure>
{structure}
</reference_structure>

<reference_copy>
{copy}
</reference_copy>

<included_phrases>
{included_phrases}
</included_phrases>

<excluded_phrases>
{excluded_phrases}
</excluded_phrases>
</product_information>"""


copy_prompt = """Generate {nversions} versions of the product description for a product with the following information.
Write in a way that target the customer.
For each version, write at least {min_length} words and at most {max_length} words.
Make sure that the structure of each output follows the structure of the reference copy.
Make sure to use the tone of voice, rythm, cadence and style of the reference copy for each output.
Use markdown format for each output.
Make sure that the structure of each output follows the structure of the reference copy.
Do not use any of the excluded words in the output.
Try to inlcude included words in the output when relevant.
Use the relevant information from the product features and intended use in the output.
Do not hallucinate any information about the product, use only the provided key features and intended use to write about the product.
Note that the reference copy should be used for style and tone only, do not use any part of the reference copy in the output. 
Do not hallucinate information about size and weight. Write about size and weight only if it is available in the list of features.
Use creative language in each output, do not use the common ways of starting product descriptions.
Avoid common phrases and cliches such as "Step into something", "Elevate something", "Discover something", "Unleash something", "Embrace something", and similar phrases.
For each version, try to write in different style.
Put each item in the list on a separate line. Use '-' character to start each item.  
Rate the quality of each version based on the following criteria:
    - how faithful it describes the product features. 
    - how well it follows the tone of voice, rythm, cadence and style of the reference copy.
    - how well it avoid the excluded words.
    - how well it includes the included words.
    - how creative the language is
The score should be a number between 0 and 10 with 10 being the best quality.
Return the result in the following JSON format:
{{
    "versions": [
        {{
            "id": 1,
            "content": The first product description,
            "explanation": A less than 20 word explanation of the score of the first product description,
            "score": The score of the first product description
        }},
        {{
            "id": 2,
            "content": The second product description,
            "explanation": A less than 20 word explanation of the score of the first product description,
            "score": The score of the second product description
        }},
        ...
    ],
    "best_version": {{
        "explanation": Explanation for why this version is the best,                                                              
        "id": The id of the best version
    }}
}}
Make sure that the output is in JSON format, no extra text should be included in the output.

<product_information>
<key_features>
{key_features}
</key_features>

<reference_copy>
{copy}
</reference_copy>

<included_phrases>
{included_phrases}
</included_phrases>

<excluded_phrases>
{excluded_phrases}
</excluded_phrases>
</product_information>"""


struct_prompt = """Generate {nversions} versions of the product description for a product with the following information.
Write in a way that target the customer.  
For each version, write at least {min_length} words and at most {max_length} words.
Make sure that the structure of each output follows the reference structure. Do not add any additional sentences or structures that are not in the reference structure.
Use markdown format for each output.
Do not include any part of the reference structure in the output.
Make sure that the structure of each output follows the reference structure.
Do not use any of the excluded words in the output.
Try to inlcude included words in the output when relevant.
Use the relevant information from the product features and intended use in the output.
Do not hallucinate any information about the product, use only the provided key features and intended use to write about the product.
Note that the reference structure should be used for structure only, do not use any part of the reference structure in the output. 
Do not hallucinate information about size and weight. Write about size and weight only if it is available in the list of features.
Use creative language in each output, do not use the common ways of starting product descriptions.
Avoid common phrases and cliches such as "Step into something", "Elevate something", "Discover something", "Unleash something", "Embrace something", and similar phrases.
For each version, try to write in different style.
Put each item in the list on a separate line. Use '-' character to start each item.  
Rate the quality of each version based on the following criteria:
    - how faithful it describes the product features. 
    - how well it follows the reference structure. 
    - how well it avoid the excluded words.                                 
    - how well it includes the included words.                                 
    - how creative the language is.                                    
The score should be a number between 0 and 10 with 10 being the best quality.                              
Return the result in the following JSON format:
{{
    "versions": [                                 
        {{
            "id": 1,
            "content": The first product description,
            "explanation": A less than 20 word explanation of the score of the first product description,                         
            "score": The score of the first product description
        }},
        {{
            "id": 2,
            "content": The second product description,
            "explanation": A less than 20 word explanation of the score of the first product description,                                     
            "score": The score of the second product description
        }},
        ...                                 
    ],                                      
    "best_version": {{
        "explanation": Explanation for why this version is the best,                                                              
        "id": The id of the best version                                                                
    }}                                                          
}}                              
Make sure that the output is in JSON format, no extra text should be included in the output.                                
                                 
<product_information>
<key_features>
{key_features}
</key_features>

<reference_structure>
{structure}
</reference_structure>

<included_phrases>
{included_phrases}
</included_phrases>

<excluded_phrases>
{excluded_phrases}
</excluded_phrases>
</product_information>"""

evaluation_prompt = """You will be given information of a product and a list of product descriptions.
Evaluate the quality of the product descriptions based on the following criteria:
    - how faithful it describes the product features. 
    - how well it follows the reference structure. 
    - how well it follows the tone of voice, rythm, cadence and style of the reference copy.
    - how well it avoid the excluded words.                                 
    - how well it includes the included words.                                 
    - how creative the language is.    
Give a score between 0 and 10 for each product description based on the above criteria.
Return the result in the following JSON format:
{{
    "versions": [                                 
        {{
            "id": 1,
            "content": The first product description,
            "explanation": A less than 20 word explanation of the score of the first product description,                         
            "score": The score of the first product description
        }},
        {{
            "id": 2,
            "content": The second product description,
            "explanation": A less than 20 word explanation of the score of the first product description,                                     
            "score": The score of the second product description
        }},
        ...                                 
    ],                                      
    "best_version": {{
        "explanation": Explanation for why this version is the best,                                                              
        "id": The id of the best version                                                                
    }}                                                          
}}                              
Make sure that the output is in JSON format, no extra text should be included in the output.  

<product_information>
<key_features>
{key_features}
</key_features>

<reference_structure>
{structure}
</reference_structure>

<reference_copy>
{copy}
</reference_copy>

<included_phrases>
{included_phrases}
</included_phrases>

<excluded_phrases>
{excluded_phrases}
</excluded_phrases>
</product_information>

<product_descriptions>
{product_descriptions}
</product_descriptions>"""


improve_structure_prompt = """You are given a structure for a product description. 
Reformat the structure so that it is easier to read and understand for human and AI. 
Return the reformatted structure only. Do not add any preceding or trailing characters.

<structure>/n{structure}</structure>"""

gemini_prompt = """
You are given information of a product, a reference structure, and a reference copy.
Please analyze the structure, make a plan on how to follow the structure correctly.
Use the tone of voice of the reference copy if available for the generated description.
Do not hallucinate, do not add information that is not in the product information.
Try your best to avoid using the excluded words and phrases.
Try your best to include the included words and phrases in the product description.
Make sure to keep to the requirements of the structure, writing no more or less than specified.
Do not enclose the output in html tags, quotes, braces, brackets or anything.
Return the product description only.

<product_information>
{key_features}
</product_information>

<reference_structure>
{structure}
</reference_structure>

<reference_copy>
{copy}
</reference_copy>

<excluded_phrases>
{excluded_phrases}
</excluded_phrases>

<included_phrases>
{included_phrases}
</included_phrases>"""


import base64
import requests
from PIL import Image
import io

# OpenAI API Key

# Function to encode the image
# def encode_image(image_path):
#   with open(image_path, "rb") as image_file:
#     return base64.b64encode(image_file.read()).decode('utf-8')


def encode_image(image_path, img_size=512):
    with open(image_path, "rb") as image_file:
        # b64_img = base64.b64encode(image_file.read()).decode('utf-8')
        # b64_img = base64.b64decode(base64.b64encode(response.content))
        img = Image.open(image_file)
        w, h = img.size 
        max_size = max(w, h)
        if max_size > img_size:
            ratio = img_size / max_size
            print("Original size: ", w, h)
            w, h = int(w * ratio), int(h * ratio)
            print("Resizing image to ", w, h)
            img = img.resize(size=(w, h))
        buffer = io.BytesIO()
        img.save(buffer, format="PNG")
        return base64.b64encode(buffer.getvalue()).decode('utf-8')


import json


def get_json(text: str):
    text = text.strip().replace('`', '').replace('json', '')
    if text.startswith("No garment detected"):
        return {
            "features": [],
            "intended_use": [],
            "alt_text": []
        }
    return json.loads(text)


@weave.op
def detect_features(image_paths, garment_type, language="English"):
    # Path to your image
    # image_path = "path_to_your_image.jpg"

    # Getting the base64 string
    try:
        base64_images = [encode_image(image_path[0]) for image_path in image_paths]
        if garment_type == "" or garment_type == "all":
            garment_type = "garment"
    
        headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {openai.api_key}"
        }

        payload = {
        "model": "gpt-4o",
        # "model": "gpt-4-turbo",
        "messages": [
            {
            "role": "user",
            "content": [
                {
                "type": "text",
                "text": f"""Describe the features of the {garment_type} in the photos.
                Return the result in in the following JSON format without any preceding or trailing text: 

                {{
                    "features": [list of comma separated features], 
                    "intended_use": [list of comma separated intended uses], 
                    "alt_text": [list of alt text for image 1, alt text for image 2]
                }}"""
                },
            ] + [{
                "type": "image_url",
                "detail": "low",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{base64_image}"
                }
                } for base64_image in base64_images]
            }
        ],
        "temperature": 0.0,
        "max_tokens": 300
        }

        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)

        print(response)

        response = response.json()
        print("image features", response["choices"][0]['message']['content'])
        jresponse = get_json(response["choices"][0]['message']['content'])
        return jresponse, base64_images
    except Exception as e:
        print(e.__class__, e)
        traceback.print_exc()
        return "", []
    

import re
def parse_structure(struct_ref):
    languages = ["_"] * (len(struct_ref) // 2)
    types = ["_"] * (len(struct_ref) // 2)
    for si in range(0, len(struct_ref), 2):
        parts = re.findall('[a-zA-Z\n ]+', struct_ref[si])
        for idx, part in enumerate(parts):
            if "language" in part.lower():
                lang = parts[idx + 1].strip()
                languages[si // 2] = lang
            if "type" in part.lower():
                type = parts[idx + 1].strip()
                types[si // 2] = type

    return types, languages


@weave.op
def detect_language(texts, model):
    langs = ["_"] * len(texts)

    try:
        messages = []
        lang_map = {}
        for i, text in enumerate(texts):
            if len(text.strip()) > 0:
                lang_mess = [HumanMessage(content=f"What is the language of the following text? Output the language only. "
                                        f"\n ```{text}```")]
                print(f"{lang_mess=}")
                messages.append(lang_mess)
                lang_map[i] = len(messages) - 1 
        detected_langs = model.batch(messages)
        print(f"{detected_langs=}")
        for k, v in lang_map.items():
            langs[k] = detected_langs[v].content
    except Exception as e:
        print(e.__class__, e)
        traceback.print_exc()

    return langs


def get_language(struct_lang, copy_lang):
    if struct_lang != "_":
        return struct_lang
    if copy_lang != "_":
        return copy_lang
    return "English"


def post_process(text: str, guidance_prompt: str, language: str, chat: ChatOpenAI):
    if guidance_prompt.strip() == "":
        return text
    
    messages = [
        SystemMessage(content=f"""You are a helpful assistant that edit documents based on the guidelines provided. 
Make sure to write in {language} language."""),
        HumanMessage(content=f"""Given the following product description, your task is to
make minimal modification to the product description such that the resulting description 
follows the rules defined in the guidelines. Make sure to preserve the structure of the 
original text as much as possible. Do not modify the structure of the original text.
Do not change the language of the original text.                     
Output only the modified text in markdown format.
Do not add any preceding or trailing text to the output.

Product description: 
{text}

Guidelines: 
{guidance_prompt}""")
    ]

    if chat is None:
        chat = ChatAnthropic(model_name="claude-3-7-sonnet-latest",
                             anthropic_api_key=os.environ["ANTHROPIC_API_KEY"],
                             max_tokens_to_sample=4096,
                             temperature=0.0)
    response = chat.invoke(messages, temperature=0.0)
    text = response.content
    return text  


def get_model(model_name, temperature=0.0):
    if model_name.startswith("gpt"):
        chat = ChatOpenAI(model=model_name, max_tokens=4096, temperature=temperature)
    elif model_name.startswith("claude"):
        chat = ChatAnthropic(model_name=model_name, 
                             anthropic_api_key=os.environ["ANTHROPIC_API_KEY"], 
                             max_tokens_to_sample=4096, 
                             temperature=temperature)
    elif model_name.startswith("gemini"):
        # chat = ChatGoogleGenerativeAI(model=model_name, 
        #                               api_key=os.environ["GOOGLE_API_KEY"])
        chat = genai.GenerativeModel(model_name)
    else:
        chat = None 
        raise ValueError(f"Model {model_name} not supported")
    return chat
    

def build_glossary(glossary_file, fieldnames=None) -> VectorStoreRetriever:
    loader = CSVLoader(file_path=glossary_file, 
                   csv_args={"delimiter": ",", 
                             "quotechar": '"'})  
                            #  "fieldnames": fieldnames})

    docs = loader.load()
    vectorstore = Chroma.from_documents(documents=docs, embedding=OpenAIEmbeddings())
    retriever = vectorstore.as_retriever()

    return retriever


def improve_structure(chat: ChatOpenAI, structure: str):
    messages = [
        HumanMessage(content=improve_structure_prompt.format(structure=structure)),]
    response = chat.invoke(messages)
    print("Original structure: ", structure)
    print("Improved structure: ", response.content)
    return response.content


def evaluate(descriptions, 
             reference_structure, 
             reference_copy, 
             key_features, 
             included_phrases, 
             excluded_phrases, 
             language, 
             chat):
    messages = [
        SystemMessage(content=f"""You are a helpful assistant that evaluates product descriptions based on the guidelines provided. Make sure to write in {language} language."""),
        HumanMessage(content=evaluation_prompt.format(key_features=key_features, 
                                                      structure=reference_structure, 
                                                      copy=reference_copy, 
                                                      included_phrases=included_phrases, 
                                                      excluded_phrases=excluded_phrases, 
                                                      product_descriptions=descriptions)),]
    
    response = chat.invoke(messages, temperature=0.0)
    print(response)
    return response


@weave.op
def generate(*data):
    global visible
    print("visible", visible)
    
    nargs = 11 
    feature, image, garment_type, model, temperature, nversions, excluded_phrases, included_phrases, glossary_upload, debug, guidance_prompt = data[:nargs]
    struct_ref = data[nargs:]

    print(f"{feature=}")
    print(f"{image=}")
    print(f"{garment_type=}")
    print(f"{model=}")
    print(f"{temperature=}")
    print(f"{excluded_phrases=}")
    print(f"{included_phrases=}")
    print(f"{debug=}")

    chat = get_model(model, temperature=temperature)

    types, struct_languages = parse_structure(struct_ref)
    copy_languages = detect_language([struct_ref[2 * i + 1] for i in range(visible + 1)], model=chat) 
    languages = [get_language(struct_lang=struct_lang, copy_lang=copy_lang) for struct_lang, copy_lang in zip(struct_languages, copy_languages)]

    # print("Struct languages--------------------------------------------\n", struct_languages)
    # print("Copy languages--------------------------------------------\n", copy_languages)
    # print("Languages--------------------------------------------\n", languages)
    # print("Types--------------------------------------------", types)

    image_features, base64_images = detect_features(image, garment_type)
    detected_features = ""
    intended_use = ""
    alt_texts = []
    key_features = feature
    if image_features is not None and len(image_features) > 0:
        alt_texts = image_features["alt_text"]
        detected_features = ", ".join(image_features["features"])
        intended_use = ", ".join(image_features["intended_use"])
        print(f"Detected features: {detected_features}, Intended use: {intended_use}, Alt text: {alt_texts}")

        key_features = key_features + ", " + detected_features + "\nIntended uses: " + intended_use

    batch = []
    min_length = 0
    max_length = 150
    response = []
    for i in range(visible + 1):
        structure = struct_ref[2 * i]
        copy = struct_ref[2 * i + 1]  

        if model.startswith("gemini"):
            if len((structure + copy).strip()) > 0:
                messages = [
                    gemini_prompt.format(min_length=min_length, 
                                         max_length=max_length, 
                                         key_features=key_features, 
                                         structure=structure, 
                                         copy=copy, 
                                         included_phrases=included_phrases, 
                                         excluded_phrases=excluded_phrases)
                ]

                batch.append(messages)
                print("Gemini input: ", messages)
                ri = chat.generate_content(messages)
                print("Gemini response: ", ri)
                response.append(ri)
        else:
            if len((structure + copy).strip()) > 0:
                if len(copy.strip()) > 0 and len(structure.strip()) > 0:    
                    print('------------')
                    print("Using both copy and structure")
                    # print("Improving structure")
                    # structure = improve_structure(chat=chat, structure=structure)
                    messages = [
                        SystemMessage(content=f"""You are a helpful assistant that writes about products for ecommerce websites. Make sure to write in {languages[i]} language."""),
                        HumanMessage(content=struct_copy_prompt.format(nversions=nversions, min_length=min_length, max_length=max_length, key_features=key_features, structure=structure, copy=copy, included_phrases=included_phrases, excluded_phrases=excluded_phrases)),]
                    
                elif len(copy.strip()) > 0:
                    print('------------')
                    print("Using copy")
                    messages = [
                        SystemMessage(content=f"""You are a helpful assistant that writes about products for ecommerce websites. Make sure to write in {languages[i]} language."""),
                        HumanMessage(content=copy_prompt.format(nversions=nversions, min_length=min_length, max_length=max_length, key_features=key_features, structure=structure, copy=copy, included_phrases=included_phrases, excluded_phrases=excluded_phrases)),]

                elif len(structure.strip()) > 0:
                    print('------------')
                    print("Using structure")
                    # print("Improving structure")
                    # structure = improve_structure(chat=chat, structure=structure)
                    messages = [
                        SystemMessage(content=f"""You are a helpful assistant that writes about products for ecommerce websites. Make sure to write in {languages[i]} language."""),
                        HumanMessage(content=struct_prompt.format(nversions=nversions, min_length=min_length, max_length=max_length, key_features=key_features, structure=structure, copy=copy, included_phrases=included_phrases, excluded_phrases=excluded_phrases)),]
                print("Chat message", messages[1].content)
                print('------------')
                batch.append(messages)

            response = chat.batch(batch)

    descriptions = ""

    descriptions = []
    descriptions_post = []
        
    if model.startswith("gemini"):
        descriptions = [msg.text for msg in response]
        descriptions_post = [post_process(text=desc, 
                                          guidance_prompt=guidance_prompt, 
                                          language=languages[i], 
                                          chat=None) for i, desc in enumerate(descriptions)]
        
        alt_texts_str = '\n\n### Alt text\n\n' + '\n- ' + '\n- '.join(alt_texts) if len(alt_texts) > 0 else ""
        alt_text_dict = {k[0]: v for (k, v) in zip(image, alt_texts)} if len(alt_texts) > 0 else {}
        result_json = {"outputs": descriptions, "alt_text": alt_text_dict}
    else:
        parser = JsonOutputParser()
        jresponse = [parser.parse(msg.content) for msg in response]
        for i, jr in enumerate(jresponse):
            print(f'{jr=}')
            bestid = jr["best_version"]["id"]
            for d in jr["versions"]:
                if d["id"] == bestid:
                    bestd = d["content"] + (f"\n\nDebug info:\n\nScore: {d['score']}\n\nExplanation: {jr['best_version']['explanation']}" if debug else "")
                    bests  = d["score"]
                    break
            
            evaluated = evaluate(descriptions=jr["versions"], 
                                reference_structure=struct_ref[2 * i], 
                                reference_copy=struct_ref[2 * i + 1], 
                                key_features=key_features, 
                                included_phrases=included_phrases, 
                                excluded_phrases=excluded_phrases, 
                                language=languages[i], chat=chat)
            print(f'{evaluated=}')
        
            bestd_post = post_process(text=bestd, 
                                      guidance_prompt=guidance_prompt, 
                                      language=languages[i], chat=chat)

            descriptions.append(bestd)
            descriptions_post.append(bestd_post)

        alt_texts_str = '\n\n### Alt text\n\n' + '\n- ' + '\n- '.join(alt_texts) if len(alt_texts) > 0 else ""
        alt_text_dict = {k[0]: v for (k, v) in zip(image, alt_texts)} if len(alt_texts) > 0 else {}

        result_json = {"outputs": jresponse if debug else descriptions, "alt_text": alt_text_dict}

    md_content = "\n\n---\n\n".join(descriptions)

    

    # post_content = post_process(text=md_content, guidance_prompt=guidance_prompt, language=languages, chat=chat)

    result_md = "Original\n\n" + md_content + "\n\n---\n\nModified\n\n" + '\n'.join(descriptions_post) + "\n\n" \
                    + alt_texts_str \
                    + '\n'.join([f'![Product photo](data:image/png;base64,{base64_image} "{alt_text}")' if base64_image != "" else "" for (base64_image, alt_text) in zip(base64_images, alt_texts)])

    

    return result_md, result_json


visible = 0
def add_output_click(*struct_ref):
    global visible
    print("Adding output ", visible)
    # print(struct_ref)
    visible += 1
    structure_texts = struct_ref[::2]
    reference_texts = struct_ref[1::2]
    structures = [gr.Textbox(label=f"Structure {i}", lines=10, value=structure_texts[i], interactive=True, visible=i <= visible) for i in range(10)]
    references = [gr.Textbox(label=f"Reference copy {i}", lines=3, value=reference_texts[i], interactive=True, visible=i <= visible) for i in range(10)]

    struct_ref = [val for pair in zip(structures, references) for val in pair]
    return struct_ref


def remove_output_click(*struct_ref):
    global visible
    print("Removing output", visible)
    if visible == 0:
        return struct_ref
    
    visible -= 1
    structure_texts = struct_ref[::2]
    reference_texts = struct_ref[1::2]
    structures = [gr.Textbox(label=f"Structure {i}", lines=10, value=structure_texts[i] if i <= visible else "", interactive=True, visible=i <= visible) for i in range(10)]
    references = [gr.Textbox(label=f"Reference copy {i}", lines=3, value=reference_texts[i] if i <= visible else "", interactive=True, visible=i <= visible) for i in range(10)]

    struct_ref = [val for pair in zip(structures, references) for val in pair]
    return struct_ref


# def show_advanced(model, temperature):
#     model = gr.Dropdown(models, value="gpt-4-turbo", interactive=True, label="Model", visible=True)
#     temperature = gr.Slider(minimum=0., maximum=1.0, value=0., interactive=True, label="Temperature", visible=True)
#     return model, temperature


with gr.Blocks() as demo:
    visible = 0
    print("Building interface")
    with gr.Row():
        with gr.Column():
            feature = gr.Textbox(label="Features", value=feature_text, lines=3, interactive=True)
            image = gr.Gallery(label="Images")
            garment_type = gr.Textbox(label="Garment Type", value="all", lines=1, interactive=True) 
            # language = gr.Dropdown(languages, value="American English", interactive=True, label="Language")
            with gr.Accordion(label="Advanced Options", open=False):
                model = gr.Dropdown(models, value=default_model, interactive=True, label="Model", visible=True)
                temperature = gr.Slider(minimum=0., maximum=1.0, value=0., interactive=True, label="Temperature", visible=True)
                nversions = gr.Slider(minimum=1, maximum=10, value=5, step=int, interactive=True, label="Number of versions", visible=True)
                excluded_phrases = gr.Textbox(label="Excluded words", interactive=True, lines=2)
                included_phrases = gr.Textbox(label="Included words", interactive=True, lines=2)
                # glossary = gr.Dataframe(row_count = (2, "dynamic"), col_count=(2,"static"), headers=["Description", "Way of writing"], label="Glossary", interactive=True)
                glossary_upload = gr.UploadButton(label="Upload Glossary", interactive=True, file_types=["csv"])
                debug = gr.Checkbox(label="Debug", interactive=True, value=True)
            with gr.Row():
                submit = gr.Button(value="Submit")
                # advanced = gr.Button(value="Advanced")
            
        with gr.Column():
            visible = 0
            guidance_prompt = gr.Textbox(label="Guidance Prompt", value=guidance_prompt_text, lines=3, interactive=True)
            struct_ref = [val for i in range(10) for val in 
                          [gr.Textbox(label=f"Structure {i}", lines=10, value="", interactive=True, visible=i <= visible), 
                           gr.Textbox(label=f"Reference copy {i}", lines=3, value="", interactive=True, visible=i <= visible)]]
            struct_ref[0].value = structure_text
            struct_ref[1].value = reference_text
            # struct_ref[2].value = structure_text_2
            with gr.Row():
                add_output = gr.Button(value="Add Output")
                remove_output = gr.Button(value="Remove Output")
                add_output.click(add_output_click, inputs=struct_ref, outputs=struct_ref)
                remove_output.click(remove_output_click, inputs=struct_ref, outputs=struct_ref)
        with gr.Column():
            md_output = gr.Markdown(label="Output", show_label=True, line_breaks=True)
            json_output = gr.JSON(label="JSON Output")  
        submit.click(generate, inputs=[feature, image, garment_type, model, temperature, nversions,
                                       excluded_phrases, included_phrases, glossary_upload, debug, guidance_prompt, *struct_ref], 
                     outputs=[md_output, json_output])
        # advanced.click(show_advanced, inputs=[], outputs=[model, temperature])
        

import bcrypt
def authf(username, password):
    try:
        with open("passwords.txt", "r") as f:
            for line in f.readlines():
                u, p = line.strip().split()
                # print(u, p, password)
                if u == username and bcrypt.checkpw(password.encode('utf-8'), p.encode('utf-8')):
                    return True
    except Exception as e:
        print("Error reading password", e)
        traceback.print_exc()

    return False


if __name__ == '__main__':
    # demo.launch(server_name="0.0.0.0", auth=authf)

    weave.init("Copywriting - Hugging Face")

    with open("formatted_prompts.txt", "w") as f:
        f.writelines('"write_struct_copy": "' + str2escaped(struct_copy_prompt).strip() + '",\n')
        f.writelines('"write_copy": "' + str2escaped(copy_prompt).strip() + '",\n')
        f.writelines('"write_struct": "' + str2escaped(struct_prompt).strip() + '",\n')
    demo.launch()