Update app.py

app.py

@@ -27,7 +27,6 @@ model = AutoModelForCausalLM.from_pretrained(
     load_in_4bit=True,
     device_map='auto'
 )
-
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
@@ -38,16 +37,12 @@ def get_completion(query, model, tokenizer):
     try:
         # Move model to CUDA
         model = model.cuda()
-        
         # Ensure input is on CUDA
         inputs = tokenizer(query, return_tensors="pt").to('cuda')
-        
         with torch.no_grad():
-            outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
-        
+            outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7)
         # Move outputs to CPU before decoding
         outputs = to_cpu(outputs)
-        
         return tokenizer.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
         return f"An error occurred: {str(e)}"
@@ -58,8 +53,86 @@ def get_completion(query, model, tokenizer):
 
 @spaces.GPU()
 def code_review(code_to_analyze):
-    query = f"As a code review expert, examine the following code for potential security flaws and provide guidance on secure coding practices:\n{code_to_analyze}"
-    result = get_completion(query, model, tokenizer)
+    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:
+
+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
+
+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
+
+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
+
+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:
+
+{code_to_analyze}
+
+Provide 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.format(code_to_analyze=code_to_analyze)
+    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.")
+    
+    if start_index != -1 and end_index != -1:
+        result = full_result[start_index:end_index].strip()
+    else:
+        result = "Unable to generate a proper code review. Please try again."
+    
     return result
 
 # Create Gradio interface
@@ -68,7 +141,7 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=10, label="Enter code to analyze"),
     outputs=gr.Textbox(label="Code Review Result"),
     title="Code Review Expert",
-    description="This tool analyzes code for potential security flaws and provides guidance on secure coding practices."
+    description="This tool analyzes code for potential security flaws, logic vulnerabilities, and provides guidance on secure coding practices."
 )
 
 # Launch the Gradio app