sangge/BackDoorBuster
sangge/BackDoorBuster
1
0
Fork 0
Code Issues Pull Requests Projects 1 Activity

fix:增加llm结果鲁棒性

Browse Source
This commit is contained in:
tritium0041 2024-06-04 17:13:30 +08:00
parent 65820966df
commit 72901463c6
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
detection/GPTdetection.py
View File

@ -25,7 +25,7 @@ def detectGPT(content: str):
signal.signal(signal.SIGTERM, timeout_handler) signal.signal(signal.SIGTERM, timeout_handler)
signal.alarm(10) signal.alarm(10)
client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1",api_key=api_key) client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1", api_key=api_key)
text = content text = content
# client = openai.OpenAI(api_key="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") #测试用key # client = openai.OpenAI(api_key="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") #测试用key
response = client.chat.completions.create( response = client.chat.completions.create(
@ -33,8 +33,8 @@ def detectGPT(content: str):
{ {
"role": "system", "role": "system",
"content": "You are a Python code reviewer.Read the code below and identify any potential security vulnerabilities. Classify them by risk level (high, medium, low, none). Only report the line number and the risk level.\nYou should output the result as json format in one line. For example: " "content": "You are a Python code reviewer.Read the code below and identify any potential security vulnerabilities. Classify them by risk level (high, medium, low, none). Only report the line number and the risk level.\nYou should output the result as json format in one line. For example: "
'[{"Line": {the line number}, "Risk": "{choose from (high,medium,low)}","Reason":"{how it is vulnable}"}] Each of these three field is required.\n' '[{"Line": {the line number}, "Risk": "{choose from (high,medium,low)}","Reason":"{how it is vulnable}"}] Each of these three field is required.\n'
"You are required to only output the json format. Do not output any other information.\n", "You are required to only output the json format. Do not output any other information.\n",
}, },
{ {
"role": "user", "role": "user",
@ -60,7 +60,10 @@ def detectGPT(content: str):
classified_results = {"high": [], "medium": [], "low": [], "none": []} classified_results = {"high": [], "medium": [], "low": [], "none": []}
for res in res_json: for res in res_json:
classified_results[res["Risk"]].append( try:
(res["Line"], text.split("\n")[res["Line"] - 1].strip()) classified_results[res["Risk"]].append(
) (res["Line"], text.split("\n")[res["Line"] - 1].strip())
)
except IndexError:
pass
return classified_results return classified_results