fix: 添加容错代码

detection/GPTdetection.py

@@ -5,7 +5,10 @@ import openai
 
 
 def detectGPT(content: str):
-    client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+    api_key = os.getenv("OPENAI_API_KEY")
+    if api_key is None:
+        raise ValueError("env OPENAI_API_KEY no set")
+    client = openai.OpenAI(api_key=api_key)
     text = content
     # client = openai.OpenAI(api_key="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") #测试用key
     response = client.chat.completions.create(

tests/test_backdoor_detection.py

@@ -1,7 +1,9 @@
 import unittest
+import warnings
 
 from detection.backdoor_detection import find_dangerous_functions
-from detection.GPTdetection import *
+from detection.GPTdetection import detectGPT
+import os
 
 
 class TestBackdoorDetection(unittest.TestCase):
@@ -57,6 +59,9 @@ class TestBackdoorDetection(unittest.TestCase):
         )
 
     def test_gpt_risk_detection(self):
+        if os.getenv("OPENAI_API_KEY") is None:
+            warnings.warn("OPENAI_API_KEY is not set, test skipped.", UserWarning)
+            self.skipTest("OPENAI_API_KEY is not set")
         content = """import os
         os.system('rm -rf /')   # high risk
         exec('print("Hello")')  # high risk
@@ -66,6 +71,9 @@ class TestBackdoorDetection(unittest.TestCase):
         self.assertEqual(len(results["high"]), 3)
 
     def test_gpt_no_risk_detection(self):
+        if os.getenv("OPENAI_API_KEY") is None:
+            warnings.warn("OPENAI_API_KEY is not set, test skipped.", UserWarning)
+            self.skipTest("OPENAI_API_KEY is not set")
         content = """a = 10
         b = a + 5
         print('This should not be detected as risky.')
@@ -75,6 +83,11 @@ class TestBackdoorDetection(unittest.TestCase):
         self.assertEqual(len(results["medium"]), 0)
         self.assertEqual(len(results["low"]), 0)
 
+    def test_gpt_env_no_set(self):
+        content = "print('test test')"
+        with self.assertRaises(ValueError):
+            detectGPT(content)
+
 
 if __name__ == "__main__":
     unittest.main()