Update app.py

app.py

@@ -119,7 +119,7 @@ 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. Start each section with its number and title, e.g., '1. Understanding the code:'"
+    query = few_shot_prompt + "\n\n" + code_to_analyze + "\n\nProvide a detailed review of the above code following the structure in the example, 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)
     
@@ -128,14 +128,17 @@ Now, review the following code using this approach:
     processed_lines = []
     relevant_sections = ['1. Understanding the code:', '2. Potential security issues:', '3. Potential logic vulnerabilities:', 'Suggestions:']
     in_relevant_section = False
+    found_first_section = False
 
     for line in lines:
         if any(section in line for section in relevant_sections):
+            if not found_first_section:
+                found_first_section = True
             in_relevant_section = True
             processed_lines.append(line)
-        elif in_relevant_section and line.strip():
+        elif found_first_section and in_relevant_section and line.strip():
             processed_lines.append(line)
-        elif in_relevant_section and not line.strip():
+        elif found_first_section and in_relevant_section and not line.strip():
             in_relevant_section = False
 
     result = "\n".join(processed_lines).strip()