Update app.py

app.py

@@ -1,5 +1,8 @@
 import streamlit as st
 import os
+import re
+from claude import embed_base64_for_claude, create_claude_image_request_for_image_captioning, \
+    create_claude_request_for_text_completion, extract_data_from_text_xml
 from prompts import prompts
 from constants import JSON_SCHEMA_FOR_GPT, UPDATED_MODEL_ONLY_SCHEMA, JSON_SCHEMA_FOR_LOC_ONLY
 from gpt import runAssistant, checkRunStatus, retrieveThread, createAssistant, saveFileOpenAI, startAssistantThread, \
@@ -9,6 +12,9 @@ from theme import flux_generated_image, flux_generated_image_seed
 import time
 from PIL import Image
 import io
+from streamlit_gsheets import GSheetsConnection
+
+conn = st.connection("gsheets", type=GSheetsConnection)
 
 
 def process_run(st, thread_id, assistant_id):
@@ -26,6 +32,28 @@ def process_run(st, thread_id, assistant_id):
             pass
 
 
+def page5():
+    st.title('Initialize your preferences!')
+    system_prompt_passed = st.text_area("System Prompt", value=prompts["PROMPT_FOR_MOOD_AND_IDEA"],
+                                        key="System Prompt")
+    caption_system_prompt = st.text_area("Captioning System Prompt", value=prompts["CAPTION_SYSTEM_PROMPT"],
+                                         key="Caption Generation System Prompt")
+    caption_prompt = st.text_area("Caption Prompt", value=prompts["CAPTION_PROMPT"],
+                                  key="Caption Generation Prompt")
+    st.text("Running on Claude")
+    col1, col2 = st.columns([1, 2])
+    with col1:
+        if st.button("Save the Prompt"):
+            st.session_state["system_prompt"] = system_prompt_passed
+            print(st.session_state["system_prompt"])
+            st.session_state["caption_system_prompt"] = caption_system_prompt
+            st.session_state["caption_prompt"] = caption_prompt
+            st.success("Saved your prompts")
+    with col2:
+        if st.button("Start Testing!"):
+            st.session_state['page'] = "Page 1"
+
+
 def page1():
     st.title("Upload Product")
     st.markdown("<h2 style='color:#FF5733; font-weight:bold;'>Add a Product</h2>", unsafe_allow_html=True)
@@ -79,54 +107,141 @@ def page2():
 
 
 def page3():
+    import random
     st.title("Scene Suggestions")
     st.write("Based on your uploaded product and references!")
     feedback = st.chat_input("Provide feedback:")
     if not st.session_state.get("assistant_initialized", False):
-        assistant_id = createAssistant("You are a helpful assistant who is an expert in Fashion Shoots.")
-        updated_prompt = prompts["IDEA_GENERATION_PROMPT"].format(
-            brand_details=st.session_state["brand_summary"],
-            product_details=st.session_state["product_info"],
-            type_of_shoot=st.session_state["shoot_type"],
-            json_schema=JSON_SCHEMA_FOR_GPT,
-            product_name=st.session_state["product_description"]
-        )
-        file_locations = []
+        file_locations_for_product = []
         for uploaded_file in st.session_state['uploaded_files']:
             bytes_data = uploaded_file.getvalue()
             image = Image.open(io.BytesIO(bytes_data))
             image.verify()
-            location = f"temp_image_{uploaded_file.name}"
+            location = f"temp_image_{random.randint(1, 100000000)}.png"
             with open(location, "wb") as f:
                 f.write(bytes_data)
-            file_locations.append(location)
+            file_locations_for_product.append(location)
             image.close()
+        file_base64_embeds_product = [embed_base64_for_claude(location) for location in file_locations_for_product]
+        caption_list_from_claude_product = []
+        for file_embeds_base64 in file_base64_embeds_product:
+            caption_from_claude = create_claude_image_request_for_image_captioning(
+                st.session_state["caption_system_prompt"], st.session_state["caption_prompt"], file_embeds_base64)
+            caption_list_from_claude_product.append(caption_from_claude)
+        string_caption_list_product = str(caption_list_from_claude_product)
+        file_locations_for_others = []
         for uploaded_file in st.session_state['reference_images']:
             bytes_data = uploaded_file.getvalue()
             image = Image.open(io.BytesIO(bytes_data))
             image.verify()
-            location = f"temp2_image_{uploaded_file.name}"
+            location = f"temp2_image_{random.randint(1, 1000000)}.png"
             with open(location, "wb") as f:
                 f.write(bytes_data)
-            file_locations.append(location)
+            file_locations_for_others.append(location)
             image.close()
-        file_ids = [saveFileOpenAI(location) for location in file_locations]
-        thread_id = startAssistantThread(file_ids, updated_prompt, "yes", "yes")
-        st.session_state.assistant_id = assistant_id
-        st.session_state.thread_id = thread_id
+        file_base64_embeds = [embed_base64_for_claude(location) for location in file_locations_for_others]
         st.session_state.assistant_initialized = True
-        regenerate_images(thread_id, assistant_id)
+        caption_list_from_claude = []
+        for file_embeds_base64 in file_base64_embeds:
+            caption_from_claude = create_claude_image_request_for_image_captioning(
+                st.session_state["caption_system_prompt"], st.session_state["caption_prompt"], file_embeds_base64)
+            caption_list_from_claude.append(caption_from_claude)
+        string_caption_list = str(caption_list_from_claude)
+        st.session_state["caption_product"] = string_caption_list_product
+        st.session_state["additional_caption"] = string_caption_list
+        additional_info_param_for_prompt = f"Brand have provided reference images whose details are:" \
+                                           f"```{string_caption_list}```. Apart from this brand needs" \
+                                           f"{st.session_state['shoot_type']}"
+        product_info = str(string_caption_list_product) + st.session_state['product_info']
+        updated_prompt_for_claude = st.session_state["system_prompt"].format(
+            BRAND_DETAILS=st.session_state['brand_summary'],
+            PRODUCT_DETAILS=product_info,
+            ADDITIONAL_INFO=additional_info_param_for_prompt
+        )
+        st.session_state["updated_prompt"] = updated_prompt_for_claude
+        message_schema_for_claude = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": updated_prompt_for_claude
+                    }
+                ]
+            }
+        ]
+        response_from_claude = create_claude_request_for_text_completion(message_schema_for_claude)
+        campaign_pattern = r"<campaign_idea>(.*?)</campaign_idea>"
+        campaigns = re.findall(campaign_pattern, response_from_claude, re.DOTALL)
+        concat_prompt_list = []
+        for idx, campaign in enumerate(campaigns, start=1):
+            get_model_prompt = extract_data_from_text_xml(campaign, "model_prompt")
+            get_background_prompt = extract_data_from_text_xml(campaign, "background_prompt")
+            concat_prompt_flux = get_model_prompt + get_background_prompt
+            concat_prompt_list.append(concat_prompt_flux)
+        flux_generated_theme_image = []
+        for concat_prompt in concat_prompt_list:
+            theme_image = flux_generated_image(concat_prompt)
+            flux_generated_theme_image.append(theme_image["file_name"])
+        print(flux_generated_theme_image)
+        st.session_state["descriptions"] = concat_prompt_list
+        st.session_state["claude_context"] = response_from_claude
+        st.session_state["images"] = flux_generated_theme_image
     if feedback:
+        updated_context = st.session_state["claude_context"]
         if 'images' in st.session_state and 'descriptions' in st.session_state:
             for image_path in st.session_state['images']:
                 os.remove(image_path)
             del st.session_state['images']
             del st.session_state['descriptions']
-            del st.session_state["json_descriptions"]
-        addMessageToThread(st.session_state.thread_id, feedback)
-        regenerate_images(st.session_state.thread_id, st.session_state.assistant_id)
+            del st.session_state["claude_context"]
+        message_schema_for_claude = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": st.session_state["updated_prompt"]
+                    }
+                ]
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": updated_context}
+                ]
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": feedback
+                    }
+                ]
+            },
+
+        ]
+        response_from_claude = create_claude_request_for_text_completion(message_schema_for_claude)
+        campaign_pattern = r"<campaign_idea>(.*?)</campaign_idea>"
+        campaigns = re.findall(campaign_pattern, response_from_claude, re.DOTALL)
+        concat_prompt_list = []
+        for idx, campaign in enumerate(campaigns, start=1):
+            get_model_prompt = extract_data_from_text_xml(campaign, "model_prompt")
+            get_background_prompt = extract_data_from_text_xml(campaign, "background_prompt")
+            concat_prompt_flux = get_model_prompt + get_background_prompt
+            concat_prompt_list.append(concat_prompt_flux)
+        flux_generated_theme_image = []
+        for concat_prompt in concat_prompt_list:
+            theme_image = flux_generated_image(concat_prompt)
+            flux_generated_theme_image.append(theme_image["file_name"])
+        st.session_state["descriptions"] = concat_prompt_list
+        st.session_state["claude_context"] = response_from_claude
+        st.session_state["images"] = flux_generated_theme_image
     selected_image_index = None
-    cols = st.columns(1)
+    cols = st.columns(4)
     for i in range(len(st.session_state["images"])):
         with cols[i]:
             st.image(st.session_state.images[i], caption=st.session_state.descriptions[i], use_column_width=True)
@@ -141,62 +256,10 @@ def page3():
         st.session_state.page = "Page 2"
 
 
-def regenerate_images(thread_id, assistant_id):
-    """Helper function to generate images and descriptions."""
-    response_from_process_list = []
-    for _ in range(1):  # Assuming you generate 1 set of image/description
-        response_from_process = process_run(st, thread_id, assistant_id)
-        response_from_process_list.append(response_from_process)
-
-    summary_list = []
-    for final_response in response_from_process_list:
-        prompt_for_idea_summary = prompts["IDEA_SUMMARY_PROMPT"].format(
-            json_schema=str(final_response)
-        )
-        summary = create_chat_completion_request_open_ai_for_summary(prompt_for_idea_summary, "No")
-        summary_list.append(summary)
-
-    # Generate images based on the summaries
-    flux_generated_theme_image = []
-    for summary in summary_list:
-        theme_image = flux_generated_image(summary)
-        flux_generated_theme_image.append(theme_image["file_name"])
-
-    # Save the new images and descriptions in session state
-    st.session_state["images"] = flux_generated_theme_image
-    st.session_state["descriptions"] = summary_list
-    st.session_state["json_descriptions"] = response_from_process_list
-
-
 def page4():
     import json
-    selected_theme_text_by_user = st.session_state.json_descriptions[st.session_state.selected_image_index]
+    selected_theme_text_by_user = st.session_state.descriptions[st.session_state.selected_image_index]
     print(selected_theme_text_by_user)
-    schema_for_model_bg = {"type": "object",
-                           "properties": {
-                               "Model": {
-                                   "type": "string",
-                                   "description": "The model name or identifier."
-                               },
-                               "Background": {
-                                   "type": "string",
-                                   "description": "Description or type of the background."
-                               }},
-                           "required": ["Model", "Background"],
-                           "additionalProperties": False
-                           }
-    session_state_desp = st.session_state["json_descriptions"]
-    prompt_to_get_details = (f"You are provided with a brief of a Fashion Shoot : "
-                             f"{session_state_desp}.\n Now provide me a JSON which will"
-                             f"have two keys ```Model``` and ```Background```. Provide all detail's"
-                             f"present about model and background in the brief provided by you. Just provide a "
-                             f"natural langauge description. I will use it as description of model and "
-                             f"background needed by the brand Output JSON following the schema")
-    response_from_open_ai = create_chat_completion_request_open_ai_for_summary(prompt_to_get_details,
-                                                                               schema_name="model_bg",
-                                                                               json_schema=schema_for_model_bg,
-                                                                               json_mode="yes")
-    json_response_from_open_ai = json.loads(response_from_open_ai)
     with (st.sidebar):
         st.title(st.session_state["product_info"])
         st.write("Product Image")
@@ -213,183 +276,74 @@ def page4():
         else:
             seed_number = 0
             st.text("Thanks will take care")
-        model_preference = st.selectbox(
-            "Model Preference",
-            ("Create Own/Edit Pre-filled", "Ideas", "Upload Reference"),
-        )
-        if model_preference == "Create Own/Edit Pre-filled":
-            pre_filled_model_details = st.text_area("Model Idea", value=json_response_from_open_ai["Model"],
-                                                    key="Model Idea")
-        elif model_preference == "Ideas":
-            prompt_to_generate_idea = ("Your task is to create model ideas for shoot of a  product of a brand. "
-                                       "The details about the brand: ```{brand_details}.\n The product: {product_name},"
-                                       "which is: ```{product_details}```.\n Reference images for the product and "
-                                       "brands shoot idea is already provided with you. Additionally brand wants to "
-                                       "have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
-                                       "think step by step and provide your ideas about what type of model the brand"
-                                       "should need based on mentioned JSON format. Also provide a combined prompt "
-                                       "which the brand will use to create a shoot image. While creating the "
-                                       "combined prompt as mentioned in the JSON schema, do not miss any details you"
-                                       " mentioned in the JSON.")
-            updated_model_idea_gen_prompt = prompt_to_generate_idea.format(
-                brand_details=st.session_state["brand_summary"],
-                product_details=st.session_state["product_info"],
-                type_of_shoot=st.session_state["shoot_type"],
-                product_name=st.session_state["product_description"]
-
-            )
-            response_for_only_model = create_chat_completion_request_open_ai_for_summary(updated_model_idea_gen_prompt
-                                                                                         , schema_name="model_only",
-                                                                                         json_schema=
-                                                                                         UPDATED_MODEL_ONLY_SCHEMA,
-                                                                                         json_mode="yes")
-            pre_filled_model_details = st.text_area("Model Idea", value=response_for_only_model,
-                                                    key="Model Idea")
-        else:
-            uploaded_files = st.file_uploader("Upload one Model Reference Image here",
-                                              accept_multiple_files=False, key="uploader")
-            bytes_data = uploaded_files.getvalue()
-            image = Image.open(io.BytesIO(bytes_data))
-            image.verify()
-            location = f"temp_image_{uploaded_files.name}"
-            with open(location, "wb") as f:
-                f.write(bytes_data)
-            image.close()
-            prompt_to_generate_idea = ("Follow this JSON Schema : {json_schema_model_only}."
-                                       "Your task is to create model ideas for shoot of a  product of a brand. "
-                                       "The details about the brand: ```{brand_details}.\n The product: {product_name},"
-                                       "which is: ```{product_details}```.\n Reference images for the product and "
-                                       "brands shoot idea is already provided with you. Additionally brand wants to "
-                                       "have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
-                                       "think step by step and provide your ideas about what type of model the brand"
-                                       "should need based on mentioned JSON format. Also provide a combined prompt "
-                                       "which the brand will use to create a shoot image. While creating the "
-                                       "combined prompt as mentioned in the JSON schema, do not miss any details you"
-                                       " mentioned in the JSON.")
-            updated_model_idea_gen_prompt = prompt_to_generate_idea.format(
-                json_schema_model_only=UPDATED_MODEL_ONLY_SCHEMA,
-                brand_details=st.session_state["brand_summary"],
-                product_details=st.session_state["product_info"],
-                type_of_shoot=st.session_state["shoot_type"],
-                product_name=st.session_state["product_description"]
-
-            )
-            json_response = create_image_completion_request_gpt(location, updated_model_idea_gen_prompt)
-            pre_filled_model_details = st.text_area("Model Idea", value=json_response,
-                                                    key="Model Idea")
-        background_preference = st.selectbox(
-            "Background Preference",
-            ("Create Own/Edit Pre-filled", "Ideas", "Upload Reference"),
-        )
-        if background_preference == "Create Own/Edit Pre-filled":
-            pre_filled_background_details = st.text_area("Background Idea",
-                                                         value=json_response_from_open_ai["Background"],
-                                                         key="Background Idea")
-        elif background_preference == "Ideas":
-            prompt_to_generate_idea = ("Follow this JSON Schema : {json_schema_background_only}."
-                                       "Your task is to create location/background ideas for shoot of a  "
-                                       "product of a brand. "
-                                       "The details about the brand: ```{brand_details}.\n The product: {product_name},"
-                                       "which is: ```{product_details}```.\n Reference images for the product and "
-                                       "brands shoot idea is already provided with you. Additionally brand wants to "
-                                       "have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
-                                       "think step by step and provide your ideas about what type of location the brand"
-                                       "should need based on mentioned JSON format. Also provide a combined prompt "
-                                       "which the brand will use to create a shoot image. While creating the "
-                                       "combined prompt as mentioned in the JSON schema, do not miss any details you"
-                                       " mentioned in the JSON.")
-            updated_bg_idea_gen_prompt = prompt_to_generate_idea.format(
-                json_schema_background_only=JSON_SCHEMA_FOR_LOC_ONLY,
-                brand_details=st.session_state["brand_summary"],
-                product_details=st.session_state["product_info"],
-                type_of_shoot=st.session_state["shoot_type"],
-                product_name=st.session_state["product_description"]
-
-            )
-            response_for_only_bg = create_chat_completion_request_open_ai_for_summary(updated_bg_idea_gen_prompt,
-                                                                                      schema_name="bg_o",
-                                                                                      json_schema=JSON_SCHEMA_FOR_LOC_ONLY,
-                                                                                      json_mode="yes")
-            pre_filled_background_details = st.text_area("Background Idea", value=response_for_only_bg,
-                                                         key="Background Idea")
-        else:
-            uploaded_files = st.file_uploader("Upload one Background Reference Image here",
-                                              accept_multiple_files=False, key="uploader")
-            bytes_data = uploaded_files.getvalue()
-            image = Image.open(io.BytesIO(bytes_data))
-            image.verify()
-            location = f"temp2_image_{uploaded_files.name}"
-            with open(location, "wb") as f:
-                f.write(bytes_data)
-            image.close()
-            prompt_to_generate_idea = ("Follow this JSON Schema : {json_schema_bg_only}."
-                                       "Your task is to create Background/Location ideas for shoot of a  "
-                                       "product of a brand. "
-                                       "The details about the brand: ```{brand_details}.\n The product: {product_name},"
-                                       "which is: ```{product_details}```.\n Reference images for the product and "
-                                       "brands shoot idea is already provided with you. Additionally brand wants to "
-                                       "have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
-                                       "think step by step and provide your ideas about what type of location the brand"
-                                       "should need based on mentioned JSON format. Also provide a combined prompt "
-                                       "which the brand will use to create a shoot image. While creating the "
-                                       "combined prompt as mentioned in the JSON schema, do not miss any details you"
-                                       " mentioned in the JSON.")
-            updated_bg_idea_gen_prompt = prompt_to_generate_idea.format(
-                json_schema_bg_only=JSON_SCHEMA_FOR_LOC_ONLY,
-                brand_details=st.session_state["brand_summary"],
-                product_details=st.session_state["product_info"],
-                type_of_shoot=st.session_state["shoot_type"],
-                product_name=st.session_state["product_description"]
-
-            )
-            json_response = create_image_completion_request_gpt(location, updated_bg_idea_gen_prompt)
-            pre_filled_background_details = st.text_area("Background Idea", value=json_response,
-                                                         key="Background Idea")
+        model__bg_preference = st.text_area("Edit Model & BG Idea", value=selected_theme_text_by_user,
+                                            key="Model & BG Idea")
         start_chat = st.button("Start Chat")
     if "mood_chat_messages" not in st.session_state:
         st.session_state["mood_chat_messages"] = []
-    if seed and dimensions and model_preference and background_preference:
+    if seed and dimensions and model__bg_preference:
         if start_chat:
-            final_mood_board_image_prompt = prompts["FINAL_PROMPT_GENERATION"].format(
-                brand_details=st.session_state["brand_summary"],
-                product_details=st.session_state["product_info"],
-                type_of_shoot=st.session_state["shoot_type"],
-                product_name=st.session_state["product_description"],
-                model_details=pre_filled_model_details,
-                location_details=pre_filled_background_details,
-                theme_details=str(selected_theme_text_by_user),
-                chat_history=str(st.session_state["mood_chat_messages"])
-            )
-            prompt_for_flux_mood_board = create_chat_completion_request_open_ai_for_summary(
-                final_mood_board_image_prompt, "No", system_message=prompts["SYSTEM_PROMPT_FOR_MOOD_BOARD"])
             if seed == "Fixed":
-                generated_flux_image = flux_generated_image_seed(prompt_for_flux_mood_board, seed_number, dimensions)
+                generated_flux_image = flux_generated_image_seed(model__bg_preference, seed_number, dimensions)
             else:
-                generated_flux_image = flux_generated_image(prompt_for_flux_mood_board)
+                generated_flux_image = flux_generated_image(model__bg_preference)
             st.session_state["mood_chat_messages"].append({
                 "role": "AI",
-                "message": prompt_for_flux_mood_board,
+                "message": model__bg_preference,
                 "image": generated_flux_image["file_name"]
             })
         # for message in st.session_state["mood_chat_messages"]:
-          #  if message["role"] == "AI":
-           #     st.write(f"Caimera AI: {message['message']}")
-            #    st.image(message['image'])
-            #else:
-             #   st.write(f"**You**: {message['message']}")
+        #  if message["role"] == "AI":
+        #     st.write(f"Caimera AI: {message['message']}")
+        #    st.image(message['image'])
+        #else:
+        #   st.write(f"**You**: {message['message']}")
         user_input = st.chat_input("Type your message here...")
         if user_input:
             st.session_state["mood_chat_messages"].append({"role": "User", "message": user_input})
-            prompt_for_flux_mood_board_n = create_chat_completion_request_open_ai_for_summary(
-                user_input, "No", system_message=prompts["SYSTEM_PROMPT_FOR_MOOD_BOARD"])
+            updated_flux_prompt = prompts["PROMPT_TO_UPDATE_IDEA_OR_MOOD"].format(
+                EXISTING_MODEL_BG_PROMPT=model__bg_preference,
+                USER_INSTRUCTIONS=user_input
+            )
+            message_schema_for_claude = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": updated_flux_prompt
+                        }
+                    ]
+                },
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": str(st.session_state["mood_chat_messages"])}
+                    ]
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": user_input + "Reference of previous conversation is also added."
+                        }
+                    ]
+                },
+            ]
+            response_from_claude = create_claude_request_for_text_completion(message_schema_for_claude)
+            cleaned_prompt = extract_data_from_text_xml(response_from_claude, "updated_prompt")
             if seed == "Fixed":
-                generated_flux_image_n = flux_generated_image_seed(prompt_for_flux_mood_board_n, seed_number,
+                generated_flux_image_n = flux_generated_image_seed(cleaned_prompt, seed_number,
                                                                    dimensions)
             else:
-                generated_flux_image_n = flux_generated_image(prompt_for_flux_mood_board_n)
+                generated_flux_image_n = flux_generated_image(cleaned_prompt)
             st.session_state["mood_chat_messages"].append({
                 "role": "AI",
-                "message": prompt_for_flux_mood_board_n,
+                "message": cleaned_prompt,
+                "actual_response": response_from_claude,
                 "image": generated_flux_image_n["file_name"]
             })
         for message in st.session_state["mood_chat_messages"]:
@@ -402,9 +356,10 @@ def page4():
 
 
 if 'page' not in st.session_state:
-    st.session_state.page = "Page 1"
+    st.session_state.page = "Page 5"
 
-# Routing between pages
+if st.session_state.page == "Page 5":
+    page5()
 if st.session_state.page == "Page 1":
     page1()
 elif st.session_state.page == "Page 2":