Update app.py

app.py

@@ -119,18 +119,27 @@ 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 your response with 'Code Review:' and end it with 'End of Review.'"
+    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."
 
     full_result = get_completion(query, model, tokenizer)
     
-    # Extract only the AI's answer
-    start_index = full_result.find("Code Review:")
-    end_index = full_result.find("End of Review.")
+    # Process the output
+    lines = full_result.split('\n')
+    processed_lines = []
+    in_review = 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:
+            processed_lines.append(line)
     
-    if start_index != -1 and end_index != -1:
-        result = full_result[start_index:end_index].strip()
+    if processed_lines:
+        result = "\n".join(processed_lines).strip()
     else:
-        result = "Unable to generate a proper code review. Please try again."
+        result = full_result.strip()  # If no markers found, return the full result
     
     return result