Update app.py

app.py

@@ -786,11 +786,19 @@ def generate_synthetic_outputs(
     # Format output variable information for the prompt
     output_vars_text = "\n".join(
         [
-            f"Key: {var['name']}: {var['description']} (Type: {var['type']}) {'Options: '+str(var['options']) if var.get('options') else ''}"
+            f"- {var['name']}: {var['description']} (Type: {var['type']}) {'Options: '+str(var['options']) if var.get('options') else ''}"
             for var in output_vars
         ]
     )
 
+    input_vars = template_spec["input"]
+    input_vars_text = "\n".join(
+        [
+            f"- {var['name']}: {var['description']} (Type: {var['type']}) {'Options: '+str(var['options']) if var.get('options') else ''}"
+            for var in input_vars
+        ]
+    )
+
     output_format = "{"
     for var in output_vars:
         output_format += f'"{var["name"]}": output, '
@@ -816,6 +824,9 @@ def generate_synthetic_outputs(
             generation_prompt = f"""
 You are generating synthetic output data based on the following input:
 
+DEFINITION OF INPUT VARIABLES:
+{input_vars_text}
+
 INPUT DATA:
 {json.dumps(input_item, indent=2)}
 
@@ -833,7 +844,7 @@ The response must be valid JSON that can be parsed directly.
 """
 
             output_data = None
-
+            print(generation_prompt)
             for attempt in range(max_retries):
                 try:
                     response = client.chat.completions.create(
@@ -939,6 +950,7 @@ The response must be valid JSON that can be parsed directly.
 
     return results
 
+
 def suggest_variable_values_from_kb(
     variable_name, variable_type, knowledge_base, client, model="gpt-3.5-turbo"
 ):
@@ -1955,10 +1967,39 @@ with tab2:
                         )
 
                     # Call LLM with the filled prompt
+                    # Create a single input data item from user inputs
+                    input_data = [st.session_state.user_inputs.copy()]
+
+                    # Create a copy of the template spec
+                    template_spec_copy = st.session_state.template_spec.copy()
+
+                    # Call generate_synthetic_outputs with the input data
                     with st.spinner("Generating output..."):
                         model_selected = st.session_state.model
-                        generated_output = call_llm(filled_prompt, model=model_selected)
-                        st.session_state.generated_output = generated_output
+                        generated_outputs = generate_synthetic_outputs(
+                            template_spec_copy,
+                            input_data,
+                            st.session_state.knowledge_base,
+                            max_retries=3,
+                        )
+
+                        # Extract the first output (since we only have one input)
+                        if generated_outputs and len(generated_outputs) > 0:
+                            # The output contains both input and output fields
+                            # We only want to display the output fields
+                            output_vars = [
+                                var["name"] for var in template_spec_copy["output"]
+                            ]
+                            output_data = {
+                                k: v
+                                for k, v in generated_outputs[0].items()
+                                if k in output_vars
+                            }
+                            st.session_state.generated_output = output_data
+                        else:
+                            st.session_state.generated_output = {
+                                "error": "Failed to generate output"
+                            }
 
             # Display generated output
             if (