File size: 28,846 Bytes
524a99c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
567930d
 
 
524a99c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b55e03e
 
 
 
 
 
 
 
524a99c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
import os, yaml
import streamlit as st
from PIL import Image
from itertools import chain

from vouchervision.model_maps import ModelMaps
from vouchervision.utils_hf import check_prompt_yaml_filename

st.set_page_config(layout="wide", page_icon='img/icon.ico', page_title='VV Prompt Builder',initial_sidebar_state="collapsed")

def create_download_button_yaml(file_path, selected_yaml_file, key_val):
    file_label = f"Download {selected_yaml_file}"
    with open(file_path, 'rb') as f:
        st.download_button(
            label=file_label,
            data=f,
            file_name=os.path.basename(file_path),
            mime='application/x-yaml',use_container_width=True,key=key_val,
        )


def upload_local_prompt_to_server(dir_prompt):
    uploaded_file = st.file_uploader("Upload a custom prompt file", type=['yaml'])
    if uploaded_file is not None:
        # Check the file extension
        file_name = uploaded_file.name
        if file_name.endswith('.yaml'):
            file_path = os.path.join(dir_prompt, file_name)
            
            # Save the file
            with open(file_path, 'wb') as f:
                f.write(uploaded_file.getbuffer())
            st.success(f"Saved file {file_name} in {dir_prompt}")
        else:
            st.error("Please upload a .yaml file that you previously created using this Prompt Builder tool.")


def save_prompt_yaml(filename, col):
    yaml_content = {
        'prompt_author': st.session_state['prompt_author'],
        'prompt_author_institution': st.session_state['prompt_author_institution'],
        'prompt_name': st.session_state['prompt_name'],
        'prompt_version': st.session_state['prompt_version'],
        'prompt_description': st.session_state['prompt_description'],
        'LLM': st.session_state['LLM'],
        'instructions': st.session_state['instructions'],
        'json_formatting_instructions': st.session_state['json_formatting_instructions'],
        'rules': st.session_state['rules'],
        'mapping': st.session_state['mapping'],
    }
    
    dir_prompt = os.path.join(st.session_state.dir_home, 'custom_prompts')
    filepath = os.path.join(dir_prompt, f"{filename}.yaml")

    with open(filepath, 'w') as file:
        yaml.safe_dump(dict(yaml_content), file, sort_keys=False)

    st.success(f"Prompt saved as '{filename}.yaml'.")

    with col: # added
        create_download_button_yaml(filepath, filename,key_val=2456237465) # added


def load_prompt_yaml(filename):
    st.session_state['user_clicked_load_prompt_yaml'] = filename
    with open(filename, 'r') as file:
        st.session_state['prompt_info'] = yaml.safe_load(file)
        st.session_state['prompt_author'] = st.session_state['prompt_info'].get('prompt_author', st.session_state['default_prompt_author']) 
        st.session_state['prompt_author_institution'] = st.session_state['prompt_info'].get('prompt_author_institution', st.session_state['default_prompt_author_institution']) 
        st.session_state['prompt_name'] = st.session_state['prompt_info'].get('prompt_name', st.session_state['default_prompt_name']) 
        st.session_state['prompt_version'] = st.session_state['prompt_info'].get('prompt_version', st.session_state['default_prompt_version']) 
        st.session_state['prompt_description'] = st.session_state['prompt_info'].get('prompt_description', st.session_state['default_prompt_description']) 
        st.session_state['instructions'] = st.session_state['prompt_info'].get('instructions', st.session_state['default_instructions']) 
        st.session_state['json_formatting_instructions'] = st.session_state['prompt_info'].get('json_formatting_instructions', st.session_state['default_json_formatting_instructions'] )
        st.session_state['rules'] = st.session_state['prompt_info'].get('rules', {})
        st.session_state['mapping'] = st.session_state['prompt_info'].get('mapping', {})
        st.session_state['LLM'] = st.session_state['prompt_info'].get('LLM', 'General Purpose')

        # print(st.session_state['mapping'].values())
        # print(chain.from_iterable(st.session_state['mapping'].values()))
        # print(list(chain.from_iterable(st.session_state['mapping'].values())))
        st.session_state['assigned_columns'] = list(chain.from_iterable(st.session_state['mapping'].values())) 


def btn_load_prompt(selected_yaml_file, dir_prompt):
    if selected_yaml_file:
        yaml_file_path = os.path.join(dir_prompt, selected_yaml_file)
        load_prompt_yaml(yaml_file_path)
    elif not selected_yaml_file:
        # Directly assigning default values since no file is selected
        st.session_state['prompt_info'] = {}
        st.session_state['prompt_author'] = st.session_state['default_prompt_author']
        st.session_state['prompt_author_institution'] = st.session_state['default_prompt_author_institution']
        st.session_state['prompt_name'] = st.session_state['prompt_name']
        st.session_state['prompt_version'] = st.session_state['prompt_version']
        st.session_state['prompt_description'] = st.session_state['default_prompt_description']
        st.session_state['instructions'] = st.session_state['default_instructions']
        st.session_state['json_formatting_instructions'] = st.session_state['default_json_formatting_instructions'] 
        st.session_state['rules'] = {}
        st.session_state['LLM'] = 'General Purpose'
        
        st.session_state['assigned_columns'] = []

        st.session_state['prompt_info'] = {
            'prompt_author': st.session_state['prompt_author'],
            'prompt_author_institution': st.session_state['prompt_author_institution'],
            'prompt_name': st.session_state['prompt_name'],
            'prompt_version': st.session_state['prompt_version'],
            'prompt_description': st.session_state['prompt_description'],
            'instructions': st.session_state['instructions'],
            'json_formatting_instructions': st.session_state['json_formatting_instructions'],
            'rules': st.session_state['rules'],
            'mapping': st.session_state['mapping'],
            'LLM': st.session_state['LLM']
        }


def check_unique_mapping_assignments():
    print(st.session_state['assigned_columns'])
    if len(st.session_state['assigned_columns']) != len(set(st.session_state['assigned_columns'])):
        st.error("Each column name must be assigned to only one category.")
        return False
    elif not st.session_state['assigned_columns']:
        st.error("No columns have been mapped.")
        return False
    elif  len(st.session_state['assigned_columns']) != len(st.session_state['rules'].keys()):
        incomplete = [item for item in list(st.session_state['rules'].keys()) if item not in st.session_state['assigned_columns']]
        st.warning(f"These columns have been mapped: {st.session_state['assigned_columns']}")
        st.error(f"However, these columns must be mapped before the prompt is complete: {incomplete}")
        return False
    else:
        st.success("Mapping confirmed.")
        return True
    

def build_LLM_prompt_config():
    col_main1, col_main2 = st.columns([10,2])
    with col_main1:
        st.session_state.logo_path = os.path.join(st.session_state.dir_home, 'img','logo.png')
        st.session_state.logo = Image.open(st.session_state.logo_path)
        st.image(st.session_state.logo, width=250)
    with col_main2:
        try:
            st.page_link('app.py', label="Home", icon="๐Ÿ ")
            st.page_link(os.path.join("pages","faqs.py"), label="FAQs", icon="โ”")
            st.page_link(os.path.join("pages","report_bugs.py"), label="Report a Bug", icon="โš ๏ธ")
        except:
            st.page_link(os.path.join(os.path.dirname(os.path.dirname(__file__)),'app.py'), label="Home", icon="๐Ÿ ")
            st.page_link(os.path.join(os.path.dirname(os.path.dirname(__file__)),"pages","faqs.py"), label="FAQs", icon="โ”")
            st.page_link(os.path.join(os.path.dirname(os.path.dirname(__file__)),"pages","report_bugs.py"), label="Report a Bug", icon="โš ๏ธ")

    st.session_state['assigned_columns'] = []
    st.session_state['default_prompt_author'] = 'unknown'
    st.session_state['default_prompt_author_institution'] = 'unknown'
    st.session_state['default_prompt_name'] = 'custom_prompt'
    st.session_state['default_prompt_version'] = 'v-1-0'
    st.session_state['default_prompt_author_institution'] = 'unknown'
    st.session_state['default_prompt_description'] = 'unknown'
    st.session_state['default_LLM'] = 'General Purpose'
    st.session_state['default_instructions'] = """1. Refactor the unstructured OCR text into a dictionary based on the JSON structure outlined below.
2. Map the unstructured OCR text to the appropriate JSON key and populate the field given the user-defined rules.
3. JSON key values are permitted to remain empty strings if the corresponding information is not found in the unstructured OCR text.
4. Duplicate dictionary fields are not allowed.
5. Ensure all JSON keys are in camel case.
6. Ensure new JSON field values follow sentence case capitalization.
7. Ensure all key-value pairs in the JSON dictionary strictly adhere to the format and data types specified in the template.
8. Ensure output JSON string is valid JSON format. It should not have trailing commas or unquoted keys.
9. Only return a JSON dictionary represented as a string. You should not explain your answer."""
    st.session_state['default_json_formatting_instructions'] = """This section provides rules for formatting each JSON value organized by the JSON key."""
    
    # Start building the Streamlit app
    col_prompt_main_left, ___, col_prompt_main_right = st.columns([6,1,3])

    
    with col_prompt_main_left:
        
        st.title("Custom LLM Prompt Builder")
        st.subheader('About')
        st.write("This form allows you to craft a prompt for your specific task. You can also edit the JSON yaml files directly, but please try loading the prompt back into this form to ensure that the formatting is correct. If this form cannot load your manually edited JSON yaml file, then it will not work in VoucherVision.")
        st.subheader(':rainbow[How it Works]')
        st.write("1. Edit this page until you are happy with your instructions. We recommend looking at the basic structure, writing down your prompt inforamtion in a Word document so that it does not randomly disappear, and then copying and pasting that info into this form once your whole prompt structure is defined.")
        st.write("2. After you enter all of your prompt instructions, click 'Save' and give your file a name.")
        st.write("3. This file will be saved as a yaml configuration file in the `..VoucherVision/custom_prompts` folder.")
        st.write("4. When you go back the main VoucherVision page you will now see your custom prompt available in the 'Prompt Version' dropdown menu.")
        st.write("5. The LLM ***only*** sees information from the 'instructions', 'rules', and 'json_formatting_instructions' sections. All other information is for versioning and integration with VoucherVisionEditor.")
        
        st.write("---")
        st.header('Load an Existing Prompt Template')
        st.write("By default, this form loads the minimum required transcription fields but does not provide rules for each field. You can also load an existing prompt as a template, editing or deleting values as needed.")

        dir_prompt = os.path.join(st.session_state.dir_home, 'custom_prompts')
        yaml_files = [f for f in os.listdir(dir_prompt) if f.endswith('.yaml')]
        col_load_text, col_load_btn, col_load_btn2 = st.columns([8,2,2])
        with col_load_text:
        # Dropdown for selecting a YAML file
            st.session_state['selected_yaml_file'] = st.selectbox('Select a prompt .YAML file to load:', [''] + yaml_files)
        with col_load_btn:
            st.write('##')
            # Button to load the selected prompt
            st.button('Load Prompt', on_click=btn_load_prompt, args=[st.session_state['selected_yaml_file'], dir_prompt],use_container_width=True)
                
        with col_load_btn2:
            if st.session_state['selected_yaml_file']:
                # Construct the full path to the file
                download_file_path = os.path.join(dir_prompt, st.session_state['selected_yaml_file'] )
                # Create the download button
                st.write('##')
                create_download_button_yaml(download_file_path, st.session_state['selected_yaml_file'],key_val=345798)

        # Prompt Author Information
        st.write("---")
        st.header("Prompt Author Information")
        st.write("We value community contributions! Please provide your name(s) (or pseudonym if you prefer) for credit. If you leave this field blank, it will say 'unknown'.")
        if 'prompt_author' not in st.session_state:# != st.session_state['default_prompt_author']:
            st.session_state['prompt_author'] = st.text_input("Enter names of prompt author(s)", value=st.session_state['default_prompt_author'],key=1111)
        else:
            st.session_state['prompt_author'] = st.text_input("Enter names of prompt author(s)", value=st.session_state['prompt_author'],key=1112)

        # Institution 
        st.write("Please provide your institution name. If you leave this field blank, it will say 'unknown'.")
        if 'prompt_author_institution' not in st.session_state:
            st.session_state['prompt_author_institution'] = st.text_input("Enter name of institution", value=st.session_state['default_prompt_author_institution'],key=1113)
        else:
            st.session_state['prompt_author_institution'] = st.text_input("Enter name of institution", value=st.session_state['prompt_author_institution'],key=1114)

        # Prompt name 
        st.write("Please provide a simple name for your prompt. If you leave this field blank, it will say 'custom_prompt'.")
        if 'prompt_name' not in st.session_state:
            st.session_state['prompt_name'] = st.text_input("Enter prompt name", value=st.session_state['default_prompt_name'],key=1115)
        else:
            st.session_state['prompt_name'] = st.text_input("Enter prompt name", value=st.session_state['prompt_name'],key=1116)
      
        # Prompt verion  
        st.write("Please provide a version identifier for your prompt. If you leave this field blank, it will say 'v-1-0'.")
        if 'prompt_version' not in st.session_state:
            st.session_state['prompt_version'] = st.text_input("Enter prompt version", value=st.session_state['default_prompt_version'],key=1117)
        else:
            st.session_state['prompt_version'] = st.text_input("Enter prompt version", value=st.session_state['prompt_version'],key=1118)
            
       
        st.write("Please provide a description of your prompt and its intended task. Is it designed for a specific collection? Taxa? Database structure?")
        if 'prompt_description' not in st.session_state:
            st.session_state['prompt_description'] = st.text_input("Enter description of prompt", value=st.session_state['default_prompt_description'],key=1119)
        else:
            st.session_state['prompt_description'] = st.text_input("Enter description of prompt", value=st.session_state['prompt_description'],key=11111)
       
        st.write('---')
        st.header("Set LLM Model Type")
        # Define the options for the dropdown
        llm_options_general = ["General Purpose",
                       "OpenAI GPT Models","Google PaLM2 Models","Google Gemini Models","MistralAI Models",]
        llm_options_all = ModelMaps.get_models_gui_list()

        if 'LLM' not in st.session_state:
            st.session_state['LLM'] = st.session_state['default_LLM']
    
        if st.session_state['LLM']:
            llm_options = llm_options_general + llm_options_all + [st.session_state['LLM']]
        else:
            llm_options = llm_options_general + llm_options_all
        # Create the dropdown and set the value to session_state['LLM']
        st.write("Which LLM is this prompt designed for? This will not restrict its use to a specific LLM, but some prompts will behave differently across models.")
        st.write("SLTPvA prompts have been validated with all supported LLMs, but perfornce may vary. If you design a prompt to work best with a specific model, then you can indicate the model here.")
        st.write("For general purpose prompts (like the SLTPvA prompts) just use the 'General Purpose' option.")
        st.session_state['LLM'] = st.selectbox('Set LLM', llm_options, index=llm_options.index(st.session_state.get('LLM', 'General Purpose')))

        st.write('---')
        # Instructions Section
        st.header("Instructions")
        st.write("These are the general instructions that guide the LLM through the transcription task. We recommend using the default instructions unless you have a specific reason to change them.")
        
        if 'instructions' not in st.session_state:
            st.session_state['instructions'] = st.text_area("Enter guiding instructions", value=st.session_state['default_instructions'].strip(), height=350,key=111112)
        else:
            st.session_state['instructions'] = st.text_area("Enter guiding instructions", value=st.session_state['instructions'].strip(), height=350,key=111112)


        st.write('---')

        # Column Instructions Section
        st.header("JSON Formatting Instructions")
        st.write("The following section tells the LLM how we want to structure the JSON dictionary. We do not recommend changing this section because it would likely result in unstable and inconsistent behavior.")
        if 'json_formatting_instructions' not in st.session_state:
            st.session_state['json_formatting_instructions'] = st.text_area("Enter general JSON guidelines", value=st.session_state['default_json_formatting_instructions'],key=111114)
        else:
            st.session_state['json_formatting_instructions'] = st.text_area("Enter general JSON guidelines", value=st.session_state['json_formatting_instructions'],key=111115)
        





        st.write('---')
        col_left, col_right = st.columns([6,4])

        null_value_rules = ''
        c_name = "EXAMPLE_COLUMN_NAME"
        c_value = "REPLACE WITH DESCRIPTION"

        with col_left:
            st.subheader('Add/Edit Columns')
            st.markdown("The pre-populated fields are REQUIRED for downstream validation steps. They must be in all prompts.")
            
            # Initialize rules in session state if not already present
            if 'rules' not in st.session_state or not st.session_state['rules']:
                for required_col in st.session_state['required_fields']:
                    st.session_state['rules'][required_col] = c_value

            


            # Layout for adding a new column name
            # col_text, col_textbtn = st.columns([8, 2])
            # with col_text:
            st.session_state['new_column_name'] = st.text_input("Enter a new column name:")
            # with col_textbtn:
            # st.write('##')
            if st.button("Add New Column") and st.session_state['new_column_name']:
                if st.session_state['new_column_name'] not in st.session_state['rules']:
                    st.session_state['rules'][st.session_state['new_column_name']] = c_value
                    st.success(f"New column '{st.session_state['new_column_name']}' added. Now you can edit its properties.")
                    st.session_state['new_column_name'] = ''
                else:
                    st.error("Column name already exists. Please enter a unique column name.")
                    st.session_state['new_column_name'] = ''
            
            # Get columns excluding the protected "catalogNumber"
            st.write('#')
            # required_columns = [col for col in st.session_state['rules'] if col not in st.session_state['required_fields']]
            editable_columns = [col for col in st.session_state['rules'] if col not in ["catalogNumber"]]
            removable_columns = [col for col in st.session_state['rules'] if col not in st.session_state['required_fields']]

            st.session_state['current_rule'] = st.selectbox("Select a column to edit:", [""] + editable_columns)
            # column_name = st.selectbox("Select a column to edit:", editable_columns)

            # Form for input fields
            with st.form(key='rule_form'):
                # format_options = ["verbatim transcription", "spell check transcription", "boolean yes no", "boolean 1 0", "integer", "[list]", "yyyy-mm-dd"]
                # current_rule["format"] = st.selectbox("Format:", format_options, index=format_options.index(current_rule["format"]) if current_rule["format"] else 0)
                # current_rule["null_value"] = st.text_input("Null value:", value=current_rule["null_value"])
                if st.session_state['current_rule']:
                    current_rule_description = st.text_area("Description of category:", value=st.session_state['rules'][st.session_state['current_rule']])
                else:
                    current_rule_description = ''
                commit_button = st.form_submit_button("Commit Column")

            # Handle commit action
            if commit_button and st.session_state['current_rule']:
                # Commit the rules to the session state.
                st.session_state['rules'][st.session_state['current_rule']] = current_rule_description
                st.success(f"Column '{st.session_state['current_rule']}' added/updated in rules.")

                # Force the form to reset by clearing the fields from the session state
                st.session_state.pop('current_rule', None)  # Clear the selected column to force reset

            delete_column_name = st.selectbox("Select a column to delete:", [""] + removable_columns)
            # with del_colbtn:
            # st.write('##')
            if st.button("Delete Column") and delete_column_name:
                del st.session_state['rules'][delete_column_name]
                st.success(f"Column '{delete_column_name}' removed from rules.")

        with col_right:
            # Display the current state of the JSON rules
            st.subheader('Formatted Columns')
            st.json(st.session_state['rules'])

        st.write('---')

        col_left_mapping, col_right_mapping = st.columns([6,4])
        with col_left_mapping:
            st.header("Mapping")
            st.write("Assign each column name to a single category.")
            st.session_state['refresh_mapping'] = False

            # Dynamically create a list of all column names that can be assigned
            # This assumes that the column names are the keys in the dictionary under 'rules'
            all_column_names = list(st.session_state['rules'].keys())

            categories = ['TAXONOMY', 'GEOGRAPHY', 'LOCALITY', 'COLLECTING', 'MISC']
            if ('mapping' not in st.session_state) or (st.session_state['mapping'] == {}):
                st.session_state['mapping'] = {category: [] for category in categories}
            for category in categories:
                # Filter out the already assigned columns
                available_columns = [col for col in all_column_names if col not in st.session_state['assigned_columns'] or col in st.session_state['mapping'].get(category, [])]

                # Ensure the current mapping is a subset of the available options
                current_mapping = [col for col in st.session_state['mapping'].get(category, []) if col in available_columns]

                # Provide a safe default if the current mapping is empty or contains invalid options
                safe_default = current_mapping if all(col in available_columns for col in current_mapping) else []

                # Create a multi-select widget for the category with a safe default
                selected_columns = st.multiselect(
                    f"Select columns for {category}:",
                    available_columns,
                    default=safe_default,
                    key=f"mapping_{category}"
                )
                # Update the assigned_columns based on the selections
                for col in current_mapping:
                    if col not in selected_columns and col in st.session_state['assigned_columns']:
                        st.session_state['assigned_columns'].remove(col)
                        st.session_state['refresh_mapping'] = True

                for col in selected_columns:
                    if col not in st.session_state['assigned_columns']:
                        st.session_state['assigned_columns'].append(col)
                        st.session_state['refresh_mapping'] = True

                # Update the mapping in session state when there's a change
                st.session_state['mapping'][category] = selected_columns
            if st.session_state['refresh_mapping']:
                st.session_state['refresh_mapping'] = False

        # Button to confirm and save the mapping configuration
        if st.button('Confirm Mapping'):
            if check_unique_mapping_assignments():
                # Proceed with further actions since the mapping is confirmed and unique
                pass

        with col_right_mapping:
            # Display the current state of the JSON rules
            st.subheader('Formatted Column Maps')
            st.json(st.session_state['mapping'])


        col_left_save, col_right_save = st.columns([6,4])
        with col_left_save:
            # Input for new file name
            new_filename = st.text_input("Enter filename to save your prompt as a configuration YAML:",placeholder='my_prompt_name')
            # Button to save the new YAML file
            if st.button('Save YAML', type='primary'):
                if new_filename:
                    if check_unique_mapping_assignments():
                        if check_prompt_yaml_filename(new_filename):
                            save_prompt_yaml(new_filename, col_left_save)
                        else:
                            st.error("File name can only contain letters, numbers, underscores, and dashes. Cannot contain spaces.")
                    else:
                        st.error("Mapping contains an error. Make sure that each column is assigned to only ***one*** category.")
                else:
                    st.error("Please enter a filename.")
        
            if st.button('Exit'):
                st.session_state.proceed_to_build_llm_prompt = False
                st.session_state.proceed_to_main = True
                st.rerun()


    with col_prompt_main_right:
        if st.session_state['user_clicked_load_prompt_yaml'] is None: # see if user has loaded a yaml to edit
            st.session_state['show_prompt_name_e'] = f"Prompt Status  :arrow_forward:  Building prompt from scratch"
            if st.session_state['prompt_name']:
                st.session_state['show_prompt_name_w'] = f"New Prompt Name  :arrow_forward:  {st.session_state['prompt_name']}.yaml"
            else:
                st.session_state['show_prompt_name_w'] = f"New Prompt Name  :arrow_forward:  [PLEASE SET NAME]"
        else:
            st.session_state['show_prompt_name_e'] = f"Prompt Status: Editing  :arrow_forward:  {st.session_state['selected_yaml_file']}"
            if st.session_state['prompt_name']:
                st.session_state['show_prompt_name_w'] = f"New Prompt Name  :arrow_forward:  {st.session_state['prompt_name']}.yaml"
            else:
                st.session_state['show_prompt_name_w'] = f"New Prompt Name  :arrow_forward:  [PLEASE SET NAME]"

        st.subheader(f'Full Prompt')
        st.write(st.session_state['show_prompt_name_e'])
        st.write(st.session_state['show_prompt_name_w'])
        st.write("---")
        st.session_state['prompt_info'] = {
            'prompt_author': st.session_state['prompt_author'],
            'prompt_author_institution': st.session_state['prompt_author_institution'],
            'prompt_name': st.session_state['prompt_name'],
            'prompt_version': st.session_state['prompt_version'],
            'prompt_description': st.session_state['prompt_description'],
            'LLM': st.session_state['LLM'],
            'instructions': st.session_state['instructions'],
            'json_formatting_instructions': st.session_state['json_formatting_instructions'],
            'rules': st.session_state['rules'],
            'mapping': st.session_state['mapping'],
        }
        st.json(st.session_state['prompt_info'])

build_LLM_prompt_config()