Update app.py

app.py

@@ -6,6 +6,7 @@ import numpy as np
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 import torch
+import re
 
 # Set up OpenAI client
 client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
@@ -43,6 +44,15 @@ if "follow_up_mode" not in st.session_state:
 if "generated_question" not in st.session_state:
     st.session_state.generated_question = None  # Stores the generated question for persistence
 
+if "code_template" not in st.session_state:
+    st.session_state.code_template = ""  # Stores the code template
+
+if "sample_test_case" not in st.session_state:
+    st.session_state.sample_test_case = ""  # Stores the sample test case
+
+if "expected_output" not in st.session_state:
+    st.session_state.expected_output = ""  # Stores the expected output
+
 if "debug_logs" not in st.session_state:
     st.session_state.debug_logs = None  # Stores debug logs for toggling
 
@@ -78,6 +88,32 @@ def generate_response(messages):
     
     return response.choices[0].message.content
 
+# Function to extract code template and sample test case from the generated question
+def extract_code_and_test_case(generated_question):
+    code_template = ""
+    sample_test_case = ""
+    expected_output = ""
+    
+    # Extract code template
+    code_match = re.search(r'```python(.*?)```', generated_question, re.DOTALL)
+    if code_match:
+        code_template = code_match.group(1).strip()
+    else:
+        # Default code template if none is found
+        code_template = "# Write your code here\n"
+
+    # Extract sample test case and expected output
+    test_case_match = re.search(r'Sample Input:\s*(.*?)\n', generated_question, re.DOTALL)
+    expected_output_match = re.search(r'Expected Output:\s*(.*?)\n', generated_question, re.DOTALL)
+    if test_case_match and expected_output_match:
+        sample_test_case = test_case_match.group(1).strip()
+        expected_output = expected_output_match.group(1).strip()
+    else:
+        sample_test_case = ""
+        expected_output = ""
+
+    return code_template, sample_test_case, expected_output
+
 # User input form for generating a new question
 with st.form(key="input_form"):
     company = st.text_input("Company", value="Google")  # Default value: Google
@@ -97,13 +133,17 @@ if generate_button:
     
     # Prepare a detailed prompt for GPT using the top question's details
     detailed_prompt = (
-        f"Transform this LeetCode question into a real-world interview scenario:\n\n"
+        f"Transform this LeetCode question into a real-world interview scenario.\n\n"
         f"**Company**: {top_question['company']}\n"
         f"**Question Name**: {top_question['questionName']}\n"
         f"**Difficulty Level**: {top_question['difficulty level']}\n"
         f"**Tags**: {top_question['Tags']}\n"
         f"**Content**: {top_question['Content']}\n"
-        f"\nPlease create a real-world interview question based on this information."
+        f"\nPlease create a real-world interview question based on this information. "
+        f"Include the following sections:\n\n"
+        f"- Problem Description\n"
+        f"- Code Template (in a Python code block)\n"
+        f"- Sample Input and Expected Output (clearly separated)\n"
     )
     
     # Generate response using OpenAI API with detailed prompt and debugging logs
@@ -112,6 +152,12 @@ if generate_button:
     # Store generated question in session state for persistence in sidebar and follow-up conversation state
     st.session_state.generated_question = response
 
+    # Extract code template and sample test case
+    code_template, sample_test_case, expected_output = extract_code_and_test_case(response)
+    st.session_state.code_template = code_template
+    st.session_state.sample_test_case = sample_test_case
+    st.session_state.expected_output = expected_output
+
     # Enable follow-up mode after generating the initial question
     st.session_state.follow_up_mode = True
 
@@ -164,23 +210,51 @@ if st.session_state.generated_question:
 else:
     st.sidebar.markdown("_No question generated yet._")
 
-# Right sidebar toggleable debug logs and code interpreter section
-with st.expander("Debug Logs (Toggle On/Off)", expanded=False):
-    if st.session_state.debug_logs:
-        st.write(st.session_state.debug_logs)
-
 st.sidebar.markdown("---")
 st.sidebar.markdown("## Python Code Interpreter")
-code_input = st.sidebar.text_area("Write your Python code here:")
+
+# Pre-fill code interpreter with code template after question generation
+if st.session_state.code_template:
+    code_input = st.sidebar.text_area("Write your Python code here:", value=st.session_state.code_template, height=300)
+else:
+    code_input = st.sidebar.text_area("Write your Python code here:", height=300)
+
 if st.sidebar.button("Run Code"):
     try:
+        # Prepare the code for execution
         exec_globals = {}
-        exec(code_input, exec_globals)  # Execute user-provided code safely within its own scope.
-        output_key = [k for k in exec_globals.keys() if k != "__builtins__"]
-        if output_key:
-            output_value = exec_globals[output_key[0]]
-            st.sidebar.success(f"Output: {output_value}")
+        # Create a function wrapper to execute the user's code
+        exec(f"def user_solution():\n{code_input}", exec_globals)
+        user_solution = exec_globals.get('user_solution', None)
+        
+        # Prepare sample test case execution
+        if st.session_state.sample_test_case:
+            # Assume the sample test case is in the format of arguments to the function
+            test_case = st.session_state.sample_test_case
+            # Evaluate the test case safely
+            test_args = eval(test_case)
+            if not isinstance(test_args, tuple):
+                test_args = (test_args,)
+            # Capture the output
+            returned_output = user_solution(*test_args)
+        else:
+            returned_output = user_solution()
+        
+        # Display the expected output and returned output
+        st.sidebar.markdown("### Sample Test Case Result:")
+        st.sidebar.markdown(f"**Sample Input:** {st.session_state.sample_test_case}")
+        st.sidebar.markdown(f"**Expected Output:** {st.session_state.expected_output}")
+        st.sidebar.markdown(f"**Your Output:** {returned_output}")
+        
+        # Compare outputs
+        if str(returned_output) == st.session_state.expected_output:
+            st.sidebar.success("Your output matches the expected output!")
         else:
-            st.sidebar.success("Code executed successfully!")
+            st.sidebar.error("Your output does not match the expected output.")
     except Exception as e:
         st.sidebar.error(f"Error: {e}")
+
+# Right sidebar toggleable debug logs and code interpreter section
+with st.expander("Debug Logs (Toggle On/Off)", expanded=False):
+    if st.session_state.debug_logs:
+        st.write(st.session_state.debug_logs)