Update app.py

app.py

@@ -14,7 +14,7 @@ system_prompt = (
 
 def remove_code_snippets(text):
     """
-    Removes code blocks, inline code, chain-of-thought, and debugging/step logs from the output.
+    Removes code blocks, inline code, chain-of-thought, error messages, and debugging/step logs from the output.
     """
     # Remove triple-backtick code blocks.
     text = re.sub(r"```[\s\S]+?```", "", text, flags=re.DOTALL)
@@ -22,16 +22,34 @@ def remove_code_snippets(text):
     text = re.sub(r"`[^`]+`", "", text)
     # Remove any text between <think> and </think> tags.
     text = re.sub(r"<think>[\s\S]*?</think>", "", text, flags=re.DOTALL)
-    # Remove debug/step log banners (e.g., "━━━━━ Step X ━━━━━")
-    text = re.sub(r"━+.*Step \d+.*━+", "", text)
-    # Remove any lines that start with "[Step" (which include duration and token info).
-    text = re.sub(r"\[Step \d+: Duration .*", "", text)
-    # Remove lines that mention code snippet instructions.
-    text = re.sub(r"Make sure to include code with the correct pattern.*", "", text)
-    # Finally, remove any remaining lines that seem to be debug logs.
+    # Remove block from "Make sure to include code with the correct pattern" until "<end_code>"
+    text = re.sub(r"Make sure to include code with the correct pattern[\s\S]*?<end_code>", "", text, flags=re.DOTALL)
+    
+    # Split the text into lines and filter out debug/error/step log lines.
     lines = text.splitlines()
-    cleaned_lines = [line for line in lines if not re.search(r"Step \d+|Duration", line)]
-    return "\n".join(cleaned_lines).strip()
+    filtered_lines = []
+    for line in lines:
+        # Skip lines that are clearly debug or error messages.
+        if re.search(r"Step \d+", line):
+            continue
+        if "Duration:" in line:
+            continue
+        if "Error in code parsing:" in line:
+            continue
+        if "💥 Error" in line:
+            continue
+        if "Reached max steps" in line:
+            continue
+        if "Make sure to provide correct code blobs" in line:
+            continue
+        if "Here is your code snippet:" in line:
+            continue
+        if "Code:" in line:
+            continue
+        filtered_lines.append(line)
+    
+    cleaned_text = "\n".join(filtered_lines)
+    return cleaned_text.strip()
 
 # Use only the final_answer tool.
 final_answer = FinalAnswerTool()
@@ -67,8 +85,8 @@ agent = CodeAgent(
 
 def run_agent(user_input):
     """
-    Runs the agent, then removes any internal chain-of-thought, step logs, and code snippets
-    before returning the final plain-text answer.
+    Runs the agent and then removes any internal chain-of-thought, error messages,
+    debug logs, and code snippet sections before returning the final plain-text answer.
     """
     raw_response = agent.run(user_input)
     print("Raw Agent Response:", raw_response)  # Debugging output (optional)