Update script.py

script.py

@@ -84,7 +84,7 @@ print('💪🏆🎉 Pong! Model responded successfully. Use generate("<prompt>")
 def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
     if CWE_vulnerability is None:
          ## introduce any vulnerability
-         prompt = (f"""
+         prompt = f"""
             Modify the following secure code to introduce a vulnerability.
             Secure Code:
             {secure_code}
@@ -92,11 +92,9 @@ def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
             Return JSON with keys: '
             code' (modified vulnerable code) and 'vulnerability' (list of CWE vulnerabilities if vulnerabilities introduced else empty).
             """
-        )
     else:
      ## introduce specific vulnerability
-     prompt = (
-            f"""
+     prompt = f"""
             Modify the following secure code to introduce a {CWE_vulnerability} vulnerability.
 
             Secure Code:
@@ -106,14 +104,12 @@ def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
             Return JSON with keys: 'code' (modified vulnerable code) and 'vulnerability' (list of CWE if vulnerabilities introduced else empty).
 
             """
-        )
     return _model.generate(prompt)
 
 
 
 def vulnerability_detection(input_code: str):
-    prompt = (
-        f"""
+    prompt = f"""
         Analyze the following code and detect security vulnerabilities.
 
         Code:
@@ -123,13 +119,11 @@ def vulnerability_detection(input_code: str):
         Identify potential vulnerabilities that could be introduced. Your task is to assess this script for potential vulnerabilities that could be introduced.
         Return JSON with key: 'vulnerability' (list of detected CWE vulnerabilities or an empty list if none found).
         """
-    )
     return _model.generate(prompt)
 
 
 def vulnerability_fix(insecure_code: str):
-    prompt = (
-        f"""
+    prompt = f"""
         Fix the security vulnerabilities in the following code.
 
         Vulnerable Code:
@@ -138,5 +132,4 @@ def vulnerability_fix(insecure_code: str):
         Your task is to fix the security vulnerabilities in the code.
         Return JSON with keys: 'code' (secure version) and 'vulnerability' (list of fixed CWE vulnerabilities if any else empty list).
         """
-    )
     return _model.generate(prompt)