Update app.py

app.py

@@ -50,79 +50,37 @@ def get_completion(query, model, tokenizer):
         # Move model back to CPU to free up GPU memory
         model = model.cpu()
         torch.cuda.empty_cache()
+        
 
 @spaces.GPU()
 def code_review(code_to_analyze):
-    few_shot_prompt = f"""Review the following code for security vulnerabilities, logic flaws, and potential improvements:
-
-```php
-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) {{
-        return true;
-    }}
-    return false;
-}}
-```
-
-1. Understanding of the code:
-   - This function attempts to authenticate a user by checking their username and password against a database.
-   - It establishes a database connection, constructs a SQL query with the provided credentials, and executes it.
-   - If any matching rows are found, it returns true (authenticated); otherwise, it returns false.
-
-2. Potential security issues:
-   - SQL Injection vulnerability: The username and password are directly inserted into the query without sanitization.
-   - Plaintext password storage: The code suggests that passwords are stored in plaintext in the database.
-   - Hardcoded database credentials: Connection details are hardcoded, which is a security risk.
-
-3. Potential logic vulnerabilities:
-   - Multiple user authentication: The function returns true if more than one row is returned, which could lead to authentication issues if multiple users have the same credentials.
-   - No input validation: There's no checking for empty or null username/password inputs.
-
-4. Suggestions for improvement:
-   - Use prepared statements to prevent SQL injection.
-   - Implement proper password hashing (e.g., using password_hash() and password_verify()).
-   - Store database credentials securely and separately from the code.
-   - Implement proper error handling and use constant-time comparison for passwords.
-   - Add input validation for username and password.
-   - Consider using a single-row fetch instead of num_rows to ensure single-user authentication.
-
-Now, review the following code using the same approach:
+    two_shot_prompt = f"""First, understand the given code:
+Analyze the purpose, functionality, input sources, output destinations, and logical flow of the code. Identify any security-sensitive operations.
+
+Now, review the following code:
+
+{code_to_analyze}
+
+Provide a brief understanding of the code.
+
+Second, correlate the context with the input code and find vulnerabilities:
+Based on your understanding of the code, identify potential security issues, logic vulnerabilities, and areas for improvement. Consider common vulnerabilities, possible misuse, input handling, and use of security functions.
+
+Now, for the same code:
 
 {code_to_analyze}
 
 Provide a detailed review including:
-1. Understanding of the code
-2. Potential security issues
-3. Potential logic vulnerabilities
-4. Suggestions for improvement
+1. Potential security issues
+2. Potential logic vulnerabilities
+3. Suggestions for improvement
 
 Start each section with its number and title."""
 
-    full_response = get_completion(few_shot_prompt, model, tokenizer)
+    full_response = get_completion(two_shot_prompt, model, tokenizer)
     
-    # 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()
-        
-        # Find the second occurrence of "1. Understanding of the code"
-        first_occurrence = ai_response.find("1. Understanding of the code")
-        if first_occurrence != -1:
-            second_occurrence = ai_response.find("1. Understanding of the code", first_occurrence + 1)
-            if second_occurrence != -1:
-                ai_response = ai_response[second_occurrence:]
-            else:
-                # If we can't find a second occurrence, start from the first one
-                ai_response = ai_response[first_occurrence:]
-        
-        return ai_response
-    else:
-        return "Error: Unable to extract the AI's response. Here's the full output:\n\n" + full_response
-
+    # Return the full response without any processing
+    return full_response
 
 # Create Gradio interface
 iface = gr.Interface(