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
client = weave.init("Copywriting - Hugging Face", settings={"capture_code": False})

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",
          "chatgpt-4o-latest",
          "o3-mini",
          "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.
If the structure does not specify the length of the output, then write at least {min_length} words and at most {max_length} words.
If the structure specifies the length of the output, then write according to the specified length.
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 copy any part of the reference structure to the output.
The structure of the output should follow the reference structure.
Make sure to keep to the requirements of the structure, writing no more or less than specified.
Do not use the structure of the reference copy for the output.
Do not use any of the excluded words in the output.
Try to include included words in the output when relevant.
Use the relevant information from the product features in the output.
Do not hallucinate any information about the product, use only the provided key features to write about the product.
Note that the reference copy should be used for style and tone only.
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.
If there is a list in the output, put each item in the list on a separate line and use '-' character to start each item.
Rate the quality of each version based on the following criteria:
    - how well 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.
If the structure does not specify the length of the output, then write at least {min_length} words and at most {max_length} words.
If the structure specifies the length of the output, then write according to the specified length.
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.
Make sure to keep to the requirements of the structure, writing no more or less than specified.
Do not use any of the excluded words in the output.
Try to include included words in the output when relevant.
Use the relevant information from the product features in the output.
Do not hallucinate any information about the product, use only the provided key features 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.
If there is a list in the output, put each item in the list on a separate line and use '-' character to start each item.
Rate the quality of each version based on the following criteria:
    - how well 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.
If the structure does not specify the length of the output, then write at least {min_length} words and at most {max_length} words.
If the structure specifies the length of the output, then write according to the specified length.
Make sure that the structure of each output follows the reference structure. Do not add any additional sentences 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.
Make sure to keep to the requirements of the structure, writing no more or less than specified.
Do not use any of the excluded words in the output.
Try to include included words in the output when relevant.
Use the relevant information from the product features in the output.
Do not hallucinate any information about the product, use only the provided key features to write about the product.
Note that the reference structure should be used for structure only.
Do not copy any part of the reference structure to 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.
If there is a list in the output, put each item in the list on a separate line and use '-' character to start each item.
Rate the quality of each version based on the following criteria:
    - how well 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>"""

reasoning_prompt = """You are given:
- Product information
- A reference structure block written in a custom templating format
- Optional reference copy to match tone
- Lists of excluded and included phrases

Your task is to generate the content for the given block by strictly following its instructions.

### How the structure block works:
- Content to be written is enclosed in double curly braces: {{ ... }}
- Control logic (like loops) is enclosed in {{% ... %}} . Logic should not be printed but just used as a control structure.
- Each block may define attributes such as tone, format, and length using a pipe: {{ block_name | tone, format, length }}
- Always follow the format (e.g., paragraph, bulleted list, inline list) as defined
- Tone and length constraints must be respected
- Lists must be formatted in standard Markdown (e.g., "- Item")
- There might also be other instructions inside the blocks, you should follow these as well

### Tone 
-  Pay attention to the tone of voice instructions in the blocks.
-  If a reference copy is provided, match its tone of voice, choice of words and sentence structure. 
-  If the reference and the tone instructions give a direction in tone and structure that is not compatible, rely on the structure more.

### Rules:
-  Use only information provided in the product data. Do not hallucinate or invent facts.
-  Use as many of the included phrases as appropriate.
-  Avoid all excluded phrases.
-  Return only the generated content for the current block — no commentary or extra formatting.
-  Do not wrap the output in quotes, HTML, brackets, or any other symbols.

<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>"""


detect_feature_prompt = """Describe the features of the {garment_type} in the photo in less than 150 words. 
The known features of the garment are: 
<key_features>
{key_features}
</key_features>
If there are any conflicts between the detected features and the known features, use the known features in the output.
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\": alt text for the image,
    \"category\": the category of the garment, for example 'dress', 'shirt', 'pants', 'shoes', etc\n
}}"""


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=1024):
    print("Encoding image", image_path)
    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)


from weave.trace.context.call_context import set_tracing_enabled


@weave.op
def detect_features(image_paths, garment_type, key_features, 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"

        chat = get_model("gpt-4o", temperature=0.0)

        messages = [[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": detect_feature_prompt.format(
                                    garment_type=garment_type,
                                    key_features=key_features)
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{base64_image}"
                                }
                            }]
                    }
                ] for base64_image in base64_images]

        response = chat.batch(messages)
        
        with set_tracing_enabled(True):
            messages_no_images = [[
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": detect_feature_prompt.format(
                                        garment_type=garment_type,
                                        key_features=key_features)
                                },
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"{image_path[0]}"
                                    }
                                }]
                        }
                    ] for image_path in image_paths]
            
            call = client.create_call("detect_features_prompts", 
                            inputs={"messages": messages_no_images})
            client.finish_call(call=call, output=response)

        print("image features")
        jresponse = {}
        for resp in response:
            print(resp.content)
            print()
            jresponse_i = get_json(resp.content)
            jresponse["features"] = jresponse.get("features", []) + jresponse_i.get("features", [])
            jresponse["intended_use"] = jresponse.get("intended_use", []) + jresponse_i.get("intended_use", [])
            jresponse["alt_text"] = jresponse.get("alt_text", []) + [jresponse_i.get("alt_text", "")]

        feature_set = set(jresponse["features"])
        intended_use_set = set(jresponse["intended_use"])
        jresponse["features"] = list(feature_set)
        jresponse["intended_use"] = list(intended_use_set)
        print(jresponse)

        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)

    model = get_model("gpt-4o", temperature=0.0)

    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") \
        or model_name.startswith("chatgpt"):
        chat = ChatOpenAI(model=model_name, max_tokens=4096, temperature=temperature)
    elif model_name.startswith("o3"):
        chat = openai.OpenAI()
    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)

    with set_tracing_enabled(False):
        image_features, base64_images = detect_features(image_paths=image, 
                                                    garment_type=garment_type, 
                                                    key_features=feature)
    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 = [
                    reasoning_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.text)
        else:
            if model.startswith("o3"):
                messages = [
                    {
                        "role": "user",
                        "content": reasoning_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)
                    }
                ]

                ri = chat.chat.completions.create(
                    model="o3-mini",
                    messages=messages,
                )

                print("O3 input: ", messages)
                print("O3 response: ", ri.choices[0].message.content)
                response.append(ri.choices[0].message.content)
            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)

    if not(model.startswith("gemini") or model.startswith("o3")):
        response = chat.batch(batch)

    descriptions = ""

    descriptions = []
    descriptions_post = []
        
    if model.startswith("gemini") or model.startswith("o3"):
        descriptions = 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)
    print("-----------------------------------------------------------------")
    print("struct_copy_prompt\n", str2escaped(struct_copy_prompt))
    print()
    print("copy_prompt\n", str2escaped(copy_prompt))
    print() 
    print("struct_prompt\n", str2escaped(struct_prompt))
    print()
    print("reasoning_prompt\n", str2escaped(reasoning_prompt))
    print()
    print("detect_feature_prompt\n", str2escaped(detect_feature_prompt))
    print("-----------------------------------------------------------------")

    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()