Update app.py

app.py

@@ -53,97 +53,19 @@ def get_completion(query, model, tokenizer):
 
 @spaces.GPU()
 def code_review(code_to_analyze):
-    few_shot_prompt = """
-You are an expert in secure coding practices and software logic analysis, tasked with reviewing source code for potential security vulnerabilities and logic flaws. Your goal is to understand the code, identify security issues that could be exploited, and uncover any logic vulnerabilities that could lead to unintended behavior. Follow these steps for each code review:
+    query = f"""Review the following code for security vulnerabilities, logic flaws, and potential improvements:
 
-1. Understand the code
-   - Analyze the purpose and functionality of the code
-   - Identify input sources and output destinations
-   - Note any security-sensitive operations (e.g., authentication, data storage, network communication)
-   - Understand the logical flow and decision points in the code
+{code_to_analyze}
 
-2. Identify potential security issues
-   - Look for common vulnerabilities (e.g., injection flaws, broken authentication, sensitive data exposure)
-   - Consider how the code might be misused or exploited
-   - Evaluate the handling of user input and data validation
-   - Check for proper use of security functions and libraries
+Provide a detailed review including:
+1. Understanding of the code
+2. Potential security issues
+3. Potential logic vulnerabilities
+4. Suggestions for improvement
 
-3. Identify potential logic vulnerabilities
-   - Look for incorrect boolean logic in conditional statements
-   - Check for off-by-one errors in loops and array operations
-   - Identify potential race conditions in multi-threaded or asynchronous code
-   - Evaluate edge cases and boundary conditions
-   - Check for proper error handling and exception management
+Start each section with its number and title."""
 
-Here's an example of how to review code:
-
-Code to review:
-```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;
-}
-```
-
-Review:
-1. Understanding the code:
-   - The function attempts to authenticate a user based on a username and password.
-   - It queries a database to check if the credentials exist.
-   - Returns true if a matching user is found, false otherwise.
-
-2. Potential security issues:
-   - SQL Injection vulnerability: Username and password are directly inserted into the query.
-   - Passwords are stored and compared in plain text, which is a severe security risk.
-   - Hardcoded database credentials in the source code.
-   - Potential for timing attacks due to direct string comparison.
-
-3. Potential logic vulnerabilities:
-   - The function returns true if more than one row is returned, which could lead to authentication bypass if multiple users have the same credentials.
-   - No input validation on username or password, potentially allowing empty strings or null values.
-
-Suggestions:
-- Use prepared statements to prevent SQL injection.
-- Use password hashing instead of storing plain text passwords.
-- Store database credentials securely outside the source code.
-- Implement proper error handling and use constant-time comparison for passwords.
-- Ensure only one user can be authenticated at a time.
-- Add input validation for username and password.
-
-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 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)
-    
-    # Process the output
-    lines = full_result.split('\n')
-    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 found_first_section and in_relevant_section and line.strip():
-            processed_lines.append(line)
-        elif found_first_section and in_relevant_section and not line.strip():
-            in_relevant_section = False
-
-    result = "\n".join(processed_lines).strip()
-    
-    return result
+    return get_completion(query, model, tokenizer)
 
 # Create Gradio interface
 iface = gr.Interface(