Fix code execution and add auto-fix functionality

Author: haseeb-heaven

libs/general_utils.py

@@ -21,7 +21,8 @@ class GeneralUtils:
  def execute_code(self, compiler_mode: str):
  code_language = st.session_state.code_language
  generated_code = st.session_state.generated_code
- 
+ code_output = None
+
  if not generated_code or len(generated_code.strip()) == 0 or not code_language or len(code_language.strip()) == 0:
  st.toast("Generated code is empty. Cannot execute an empty code.", icon="❌")
  logger.error("Error in code execution: Generated code is empty.")
@@ -40,37 +41,37 @@ def execute_code(self, compiler_mode: str):
  return html_content
 
  else:
- output = self.run_code(generated_code, code_language)
+ code_output = self.run_code(generated_code, code_language)
 
  # Check for errors in code execution
- if "error" in output.lower() or "exception" in output.lower() or "SyntaxError" in output.lower() or "NameError" in output.lower():
+ if "error" in code_output.lower() or "exception" in code_output.lower() or "SyntaxError" in code_output.lower() or "NameError" in code_output.lower():
 
- logger.error(f"Error in code execution: {output}")
+ logger.error(f"Error in code execution: {code_output}")
  response = st.session_state.sequential_chain({'code_topic': generated_code})
  fixed_code = response['code_fix']
  st.code(fixed_code, language=code_language.lower())
 
  with st.expander('Message History'):
  st.info(st.session_state.memory.buffer)
  logger.warning(f"Trying to run fixed code: {fixed_code}")
- output = self.run_code(code=fixed_code, language=code_language)
- logger.warning(f"Fixed code output: {output}")
- logger.info(f"Execution Output: '{output}' and session output: '{st.session_state.code_output}'")
+ code_output = self.run_code(code=fixed_code, language=code_language)
+ logger.warning(f"Fixed code output: {code_output}")
+ logger.info(f"Execution Output: '{code_output}' and session output: '{st.session_state.code_output}'")
 
  # check for expected output
  if (st.session_state.code_output is not None and st.session_state.code_output != "" and len(st.session_state.code_output) > 0):
- if (output == st.session_state.code_output):
- st.toast("Output:\n" + output, icon="🔥")
+ if (code_output == st.session_state.code_output):
+ st.toast("Output:\n" + code_output, icon="🔥")
  else:
  st.toast("Error the expected output doesnt match the generated output:\n'" + st.session_state.code_output + "'\n", icon="❌")
- logger.info(f"Execution Output: {output}")
- return output
+ logger.info(f"Execution Output: {code_output}")
+ return code_output
 
- except Exception as e:
- st.toast("Error in code execution:", icon="❌")
+ except Exception as exception:
+ st.toast("Error in code execution: " + code_output, icon="❌")
  # Output the stack trace
- st.toast(traceback.format_exc(), icon="❌")
  logger.error(f"Error in code execution: {traceback.format_exc()}")
+ return code_output
 
  # Generate Dynamic HTML for JDoodle Compiler iFrame Embedding.
  def generate_dynamic_html(self,language, code_prompt):

script.py

@@ -66,7 +66,11 @@ def initialize_session_state():
  st.session_state.palm_langchain = None
  if "gemini_langchain" not in st.session_state:
  st.session_state.gemini_langchain = None
- 
+ if "code_fix_instructions" not in st.session_state:
+ st.session_state.code_fix_instructions = None
+ if "session_state.sequential_chain" not in st.session_state:
+ st.session_state.sequential_chain = None
+
  # Initialize session state for Vertex AI
  if "vertexai" not in st.session_state:
  st.session_state["vertexai"] = {
@@ -323,11 +327,12 @@ def main():
  # Input box for entering the prompt
  st.session_state.code_prompt = st.text_area("Enter Prompt", height=200, placeholder=placeholder,label_visibility='hidden')
 
- with st.expander("Input/Output Options"):
+ with st.expander("Input Options"):
  with st.container():
  st.session_state.code_input = st.text_input("Input (Stdin)", placeholder="Input (Stdin)", label_visibility='collapsed',value=st.session_state.code_input)
  st.session_state.code_output = st.text_input("Output (Stdout)", placeholder="Output (Stdout)", label_visibility='collapsed',value=st.session_state.code_output)
- 
+ st.session_state.code_fix_instructions = st.text_input("Fix instructions", placeholder="Fix instructions", label_visibility='collapsed',value=st.session_state.code_fix_instructions)
+
  # Set the input and output to None if the input and output is empty
  if st.session_state.code_input and st.session_state.code_output: 
  if len(st.session_state.code_input) == 0:
@@ -344,22 +349,22 @@ def main():
 
  with st.form('code_controls_form'):
  # Create columns for alignment
- file_name_col, save_code_col,generate_code_col,run_code_col = st.columns(4)
+ file_name_col, save_code_col,generate_code_col,run_code_col,fix_code_col = st.columns(5)
 
  # Input Box (for entering the file name) in the first column
  with file_name_col:
  code_file = st.text_input("File name", value="", placeholder="File name", label_visibility='collapsed')
 
  # Save Code button in the second column
  with save_code_col:
- download_code_submitted = st.form_submit_button("Download Code")
+ download_code_submitted = st.form_submit_button("Download")
  if download_code_submitted:
  file_format = "text/plain"
  st.session_state.download_link = general_utils.generate_download_link(st.session_state.generated_code, code_file,file_format,True)
 
  # Generate Code button in the third column
  with generate_code_col:
- button_label = "Generate Code" if st.session_state["vertexai"]["model_name"] == "code-bison" else "Complete Code"
+ button_label = "Generate" if st.session_state["vertexai"]["model_name"] == "code-bison" else "Complete"
  generate_submitted = st.form_submit_button(button_label)
 
  if generate_submitted:
@@ -415,9 +420,21 @@ def main():
  st.session_state.generated_code = ""
  logger.error(f"Please select a valid AI option selected '{st.session_state.ai_option}' option")
 
+ # Fix Code button in the fourth column
+ with fix_code_col:
+ fix_submitted = st.form_submit_button("Auto Fix")
+ if fix_submitted:
+ if len(st.session_state.code_fix_instructions) == 0:
+ st.toast("Missing fix instructions", icon="❌")
+ logger.warning("Missing fix instructions")
+ 
+ logger.info(f"Fixing code with instructions: {st.session_state.code_fix_instructions}")
+ st.session_state.generated_code = st.session_state.palm_langchain.fix_generated_code(st.session_state.generated_code, st.session_state.code_language,st.session_state.code_fix_instructions)
+
+
  # Run Code button in the fourth column
  with run_code_col:
- execute_submitted = st.form_submit_button("Execute Code")
+ execute_submitted = st.form_submit_button("Execute")
  if execute_submitted:
  st.session_state.output = general_utils.execute_code(st.session_state.compiler_mode)