Update app.py

app.py

@@ -53,18 +53,18 @@ def get_completion(query, model, tokenizer):
 
 @spaces.GPU()
 def code_review(code_to_analyze):
-    few_shot_prompt = """Review the following code for security vulnerabilities, logic flaws, and potential improvements:
+    few_shot_prompt = f"""Review the following code for security vulnerabilities, logic flaws, and potential improvements:
 
 ```php
-function authenticateUser($username, $password) {
+function authenticateUser($username, $password) {{
     $conn = new mysqli("localhost", "user", "password", "database");
     $query = "SELECT * FROM users WHERE username = '$username' AND password = '$password'";
     $result = $conn->query($query);
-    if ($result->num_rows > 0) {
+    if ($result->num_rows > 0) {{
         return true;
-    }
+    }}
     return false;
-}
+}}
 ```
 
 1. Understanding of the code:
@@ -103,14 +103,20 @@ Start each section with its number and title."""
 
     full_response = get_completion(few_shot_prompt, model, tokenizer)
     
-    # Extract only the part of the response after the input code
-    start_marker = f"Now, review the following code using the same approach:\n\n{code_to_analyze}\n\n"
-    start_index = full_response.find(start_marker)
-    if start_index != -1:
-        relevant_response = full_response[start_index + len(start_marker):].strip()
-        return relevant_response
+    # Find the start of the AI's response (after the input code)
+    response_start = full_response.find(code_to_analyze)
+    if response_start != -1:
+        response_start += len(code_to_analyze)
+        ai_response = full_response[response_start:].strip()
+        
+        # Remove any leading text before the first numbered section
+        first_section = ai_response.find("1.")
+        if first_section != -1:
+            ai_response = ai_response[first_section:]
+        
+        return ai_response
     else:
-        return "Error: Unable to extract the relevant part of the AI's response."
+        return "Error: Unable to extract the AI's response. Here's the full output:\n\n" + full_response
 
 # Create Gradio interface
 iface = gr.Interface(