add known data

app.py

@@ -1,128 +1,198 @@
 import json
+
 import gradio as gr
 from dotenv import load_dotenv
+from openai import BadRequestError
+
 from llm_api.exceptions import RefusalError
+
 # from llm_api.langchain_api import extract_info, follow_structure
 from llm_api.native_api import extract_info, follow_structure
-from openai import BadRequestError
-import os
 
+load_dotenv(".env", override=True)
+garment_categories = [
+    "Upper Garment",
+    "Lower Garment",
+    "Hat",
+    "Luggage",
+    "Handbag",
+    "Makeup",
+    "Chair",
+    "Carpet",
+    "Watch",
+    "Cookware",
+    "Tableware",
+    "Sleepware",
+    "BathroomItem",
+]
+schema_categories = [
+    "UpperGarmentProperties",
+    "LowerGarmentProperties",
+    "HatProperties",
+    "LuggageProperties",
+    "HandbagProperties",
+    "MakeupProperties",
+    "ChairProperties",
+    "CarpetProperties",
+    "WatchProperties",
+    "CookwareProperties",
+    "TablewareProperties",
+    "SleepwareProperties",
+    "BathroomItemProperties",
+]
 
-# Load environment variables
-attr_user = os.getenv('ATTR_USER')
-attr_pass = os.getenv('ATTR_PASS')
 
-load_dotenv('.env', override=True)
-garment_categories = ['Auto', 'Upper Garment', 'Lower Garment', 'Hat', 'Luggage', 'Handbag', 'Makeup', 'Chair', 'Carpet', 'Watch', 'Cookware', 'Tableware', 'Sleepware', 'BathroomItem']
-schema_categories = ['Auto', 'UpperGarmentProperties', 'LowerGarmentProperties', 'HatProperties', 'LuggageProperties', 'HandbagProperties', 'MakeupProperties', 'ChairProperties', 'CarpetProperties', 'WatchProperties', 'CookwareProperties', 'TablewareProperties', 'SleepwareProperties', 'BathroomItemProperties']
-
-def extract_free_fn(gallery, category, provider, schema):
+def extract_free_fn(gallery, category, provider, schema, known_data):
     img_paths = [item[0] for item in gallery]
     category_schema = schema_categories[garment_categories.index(category)]
     # write schema to file before calling extract_free_properties
     try:
-        if schema is not None and isinstance(schema, str) and 'import' not in schema.replace("'import'", '').replace("'important'", ''):
-            with open('schema_free.py', 'r') as f:
+        if (
+            schema is not None
+            and isinstance(schema, str)
+            and "import"
+            not in schema.replace("'import'", "").replace("'important'", "")
+        ):
+            with open("schema_free.py", "r") as f:
                 default_schema = f.read()
-            import_part = default_schema.split('#####')[0]
-            print('Executing freestyle schema...')
+            import_part = default_schema.split("#####")[0]
+            print("Executing freestyle schema...")
             full_schema = import_part + schema
 
             # dynamically change the last line of the schema to include the category schema
-            if category_schema in schema_categories[1:]:
-                print('Dynamic schema activated')
-                schema_lines = full_schema.split('\n')
-                schema_lines[-1] = f'    properties: {category_schema}'
-                full_schema = '\n'.join(schema_lines)
+            if category_schema in schema_categories:
+                print("Dynamic schema activated")
+                schema_lines = full_schema.split("\n")
+                schema_lines[-1] = f"    properties: {category_schema}"
+                full_schema = "\n".join(schema_lines)
 
-            with open('read_schema_free.py', 'w') as f:
+            with open("read_schema_free.py", "w") as f:
                 f.write(full_schema)
             exec(full_schema, globals())
-            print('Executed freestyle schema')
+            print("Executed freestyle schema")
         else:
-            print('SCHEMA ELSE')
-            raise Exception('Invalid schema')
+            print("SCHEMA ELSE")
+            raise Exception("Invalid schema")
     except Exception as e:
         print(e)
-        return 'Invalid schema! Please check the schema or Reload this page to restore default settings', {}, []
+        return (
+            "Invalid schema! Please check the schema or Reload this page to restore default settings",
+            {},
+            [],
+        )
 
     try:
-        status_code, free_attributes = extract_info(img_paths, provider, Garment)
+        status_code, free_attributes = extract_info(img_paths, provider, Garment, known_data=known_data)
         free_attributes = free_attributes.model_dump_json()
     except RefusalError as e:
         return e.message, {}, []
     except BadRequestError as e:
         return e.message, {}, []
 
-    return f'status: {status_code} - provider: {provider}', free_attributes, [1] * 100
+    return f"status: {status_code} - provider: {provider}", free_attributes, [1] * 100
+
 
-def follow_schema_fn(free_attributes, provider, probs, schema):
+def follow_schema_fn(free_attributes, provider, probs, schema, known_data):
     try:
-        if schema is not None and isinstance(schema, str) and 'import' not in schema.replace("'import'", '').replace("'important'", ''):
-            with open('schema_structured.py', 'r') as f:
+        if (
+            schema is not None
+            and isinstance(schema, str)
+            and "import"
+            not in schema.replace("'import'", "").replace("'important'", "")
+        ):
+            with open("schema_structured.py", "r") as f:
                 default_schema = f.read()
-            import_part = default_schema.split('#####')[0]
-            print('executing structured schema\n')
+            import_part = default_schema.split("#####")[0]
+            print("executing structured schema\n")
             full_schema = import_part + schema
-            with open('read_schema_structured.py', 'w') as f:
+            with open("read_schema_structured.py", "w") as f:
                 f.write(full_schema)
             exec(full_schema, globals())
-            print('executed structured schema')
+            print("executed structured schema")
         else:
-            print('SCHEMA ELSE')
-            raise Exception('Invalid schema')
+            print("SCHEMA ELSE")
+            raise Exception("Invalid schema")
     except:
-        return 'Invalid schema! Please check the schema or Reload this page to restore default settings', {}
+        return (
+            "Invalid schema! Please check the schema or Reload this page to restore default settings",
+            {},
+        )
 
     try:
-        status_code, schema_attributes = follow_structure(free_attributes, provider, Garment)
+        status_code, schema_attributes = follow_structure(
+            free_attributes, provider, Garment, known_data=known_data
+        )
         schema_attributes = schema_attributes.model_dump_json()
     except RefusalError as e:
         return e.message, {}
     except BadRequestError as e:
         return e.message, {}
-    
-    return f'status {status_code} - provider: {provider}', schema_attributes
 
+    return f"status {status_code} - provider: {provider}", schema_attributes
 
-with gr.Blocks(title='Internal Demo for Attribution') as demo:
+
+with gr.Blocks(title="Internal Demo for Attribution") as demo:
     with gr.Row():
         with gr.Column(scale=12):
-            gr.Markdown('''<div style="text-align: center; font-size: 24px;"><strong>Internal Demo for Attribution</strong></div>''')
+            gr.Markdown(
+                """<div style="text-align: center; font-size: 24px;"><strong>Internal Demo for Attribution</strong></div>"""
+            )
     with gr.Row():
         with gr.Column(scale=4):
-            # gr.Markdown('''<div style="text-align: left; font-size: 18px;">Upload images of a garment and click Extract Freestyle</div>''')
-            gallery = gr.Gallery(label='Images', type='filepath')
-            probs_text = gr.Textbox(label='Probabilities', placeholder='Probabilities', visible=False)
-            category = gr.Dropdown(label='Category', choices=garment_categories, value=garment_categories[0])
-            provider = gr.Dropdown(label='Provider', choices=['openai', 'anthropic'], value='anthropic')
-            button_extract = gr.Button(value='Extract Freestyle')
+            gallery = gr.Gallery(label='Upload images of your product here',type="filepath")
+            probs_text = gr.Textbox(
+                label="Probabilities", placeholder="Probabilities", visible=False
+            )
+            category = gr.Dropdown(
+                label="Category",
+                choices=garment_categories,
+                value=garment_categories[0],
+            )
+            provider = gr.Dropdown(
+                label="Provider", choices=["openai", "anthropic"], value="anthropic"
+            )
+            known_data = gr.TextArea(
+                label="Insert known attributes of the product here (Optional)", placeholder="", visible=True
+            )
+            button_extract = gr.Button(value="Extract Freestyle")
         with gr.Column(scale=4):
-            # gr.Markdown('''<div style="text-align: left; font-size: 18px;">Click Follow Schema to convert the freestyle output to predefined schema</div>''')
-            status_free_text = gr.Textbox(label='Request Status', placeholder='')
-            free_attributes = gr.JSON(label='Freestyle Attribute')
-            button_follow = gr.Button(value='Follow Schema')
+            status_free_text = gr.Textbox(label="Request Status", placeholder="")
+            free_attributes = gr.JSON(label="Freestyle Attribute")
+            button_follow = gr.Button(value="Follow Schema")
         with gr.Column(scale=4):
-            status_structured_text = gr.Textbox(label='Request Status', placeholder='')
-            schema_attributes = gr.JSON(label='Schema Attributes')
+            status_structured_text = gr.Textbox(label="Request Status", placeholder="")
+            schema_attributes = gr.JSON(label="Schema Attributes")
     with gr.Row():
-        # free_schema = gr.Dataframe(row_count=(4, 'dynamic'), label='Freestyle Schema', col_count=(4, 'dynamic'), interactive=True)
-        # schema = gr.Dataframe(row_count=(4, 'dynamic'), label='Schema', col_count=(4, 'dynamic'), interactive=True)
-        with open('schema_free.py', 'r') as f:
+        with open("schema_free.py", "r") as f:
             schema_free_text = f.read()
-        schema_free_text = schema_free_text.split('#####')[1]
-        with open('schema_structured.py', 'r') as f:
+        schema_free_text = schema_free_text.split("#####")[1]
+        with open("schema_structured.py", "r") as f:
             schema_structured_text = f.read()
-        schema_structured_text = schema_structured_text.split('#####')[1]
-        free_schema = gr.TextArea(label='Freestyle Schema', placeholder='Freestyle Schema', visible=True, value=schema_free_text, interactive=True)
-        structured_schema = gr.TextArea(label='Structured Schema', placeholder='Schema', visible=True, value=schema_structured_text, interactive=True)
-    
-    button_extract.click(extract_free_fn, [gallery, category,provider, free_schema], [status_free_text, free_attributes, probs_text])
-    button_follow.click(follow_schema_fn, [free_attributes, provider, probs_text, structured_schema], [status_structured_text, schema_attributes])
-
-auth = (attr_user, attr_pass)
+        schema_structured_text = schema_structured_text.split("#####")[1]
+        free_schema = gr.TextArea(
+            label="Freestyle Schema",
+            placeholder="Freestyle Schema",
+            visible=True,
+            value=schema_free_text,
+            interactive=True,
+        )
+        structured_schema = gr.TextArea(
+            label="Structured Schema",
+            placeholder="Schema",
+            visible=True,
+            value=schema_structured_text,
+            interactive=True,
+        )
 
-# demo.launch(server_name="0.0.0.0", server_port=7688, share=False, ssl_verify=True, ssl_certfile=None, ssl_keyfile=None, auth=None)
-# demo.launch(server_name="0.0.0.0", server_port=7860, share=False, ssl_verify=False, ssl_certfile="server.crt", ssl_keyfile="server.key", auth=None)
-demo.launch(auth=auth, debug=True)
+    button_extract.click(
+        extract_free_fn,
+        [gallery, category, provider, free_schema, known_data],
+        [status_free_text, free_attributes, probs_text],
+    )
+    button_follow.click(
+        follow_schema_fn,
+        [free_attributes, provider, probs_text, structured_schema, known_data],
+        [status_structured_text, schema_attributes],
+    )
 
+demo.launch()