Update app.py

app.py

@@ -119,27 +119,26 @@ Now, review the following code using this approach:
 """
 
     # Concatenate the prompt with the code to analyze
-    query = few_shot_prompt + "\n\n" + code_to_analyze + "\n\nProvide a detailed review following the structure above, including understanding the code, identifying potential security issues, identifying potential logic vulnerabilities, and offering specific suggestions for improvement."
+    query = few_shot_prompt + "\n\n" + code_to_analyze + "\n\nProvide a detailed review following the structure above, including understanding the code, identifying potential security issues, identifying potential logic vulnerabilities, and offering specific suggestions for improvement. Start each section with its number and title, e.g., '1. Understanding the code:'"
 
     full_result = get_completion(query, model, tokenizer)
     
     # Process the output
     lines = full_result.split('\n')
     processed_lines = []
-    in_review = False
+    relevant_sections = ['1. Understanding the code:', '2. Potential security issues:', '3. Potential logic vulnerabilities:', 'Suggestions:']
+    in_relevant_section = False
+
     for line in lines:
-        if line.strip().lower().startswith("code review:"):
-            in_review = True
-            continue
-        if line.strip().lower() == "end of review.":
-            break
-        if in_review:
+        if any(section in line for section in relevant_sections):
+            in_relevant_section = True
             processed_lines.append(line)
-    
-    if processed_lines:
-        result = "\n".join(processed_lines).strip()
-    else:
-        result = full_result.strip()  # If no markers found, return the full result
+        elif in_relevant_section and line.strip():
+            processed_lines.append(line)
+        elif in_relevant_section and not line.strip():
+            in_relevant_section = False
+
+    result = "\n".join(processed_lines).strip()
     
     return result