Skip to content

Commit a515680

Browse files
njmulsqbkingthorin
njmulsqb
authored and
kingthorin
committed
pass otherinfo to llm
Signed-off-by: Najam Ul Saqib <[email protected]>
1 parent 2134216 commit a515680

File tree

3 files changed

+147
-19
lines changed

3 files changed

+147
-19
lines changed

‎addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmAssistant.java‎

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,22 +34,45 @@ public interface LlmAssistant {
3434
"As a software architect, and based on your previous answer, generate other potential missing endpoints that are not mentioned in the OpenAPI file. For example, if there is GET /product/1, suggest DELETE /product/1 if it's not mentioned")
3535
HttpRequestList complete();
3636

37-
@SystemMessage(
38-
"You are a web application security expert reviewing potential false positives. Answer only in JSON.")
39-
@UserMessage(
40-
"Your task is to review the following finding from ZAP (Zed Attack Proxy).\n"
41-
+ "The confidence level is a pull down field which allows you to specify how confident you are in the validity of the finding : \n"
42-
+ "- 0 if it's False Positive\n"
43-
+ "- 1 if it's Low\n"
44-
+ "- 2 if it's Medium\n"
45-
+ "- 3 if it's High\n"
46-
+ "\n"
47-
+ "The alert is described as follows : {{description}}\n"
48-
+ "\n"
49-
+ "As evidence, the HTTP response contains :\n"
50-
+ "---\n"
51-
+ "{{evidence}}\n"
52-
+ "---\n"
53-
+ "Provide a short consistent explanation of the new score.\n")
37+
static final String ALERT_REVIEW_SYSTEM_MSG =
38+
"You are a web application security expert reviewing potential false positives. Answer only in JSON.";
39+
static final String ALERT_REVIEW_GOAL =
40+
"Provide a short consistent explanation of the new score.\n";
41+
static final String ALERT_REVIEW_PROMPT =
42+
"""
43+
Your task is to review the following finding from ZAP (Zed Attack Proxy).
44+
The confidence level is a pull down field which allows you to specify how confident you are in the validity of the finding:
45+
- 0 if it's False Positive
46+
- 1 if it's Low
47+
- 2 if it's Medium
48+
- 3 if it's High
49+
The alert is described as follows: {{description}}
50+
51+
As evidence, the HTTP message contains:
52+
---
53+
{{evidence}}
54+
---
55+
"""
56+
+ ALERT_REVIEW_GOAL;
57+
58+
static final String ALERT_REVIEW_OTHERINFO_PROMPT =
59+
ALERT_REVIEW_PROMPT
60+
+ """
61+
Also, here's some additional information that may be useful for you to reach your conclusion:
62+
---
63+
{{otherinfo}}
64+
---
65+
"""
66+
+ ALERT_REVIEW_GOAL;
67+
68+
@SystemMessage(ALERT_REVIEW_SYSTEM_MSG)
69+
@UserMessage(ALERT_REVIEW_PROMPT)
5470
Confidence review(@V("description") String description, @V("evidence") String evidence);
71+
72+
@SystemMessage(ALERT_REVIEW_SYSTEM_MSG)
73+
@UserMessage(ALERT_REVIEW_OTHERINFO_PROMPT)
74+
Confidence review(
75+
@V("description") String description,
76+
@V("evidence") String evidence,
77+
@V("otherinfo") String otherinfo);
5578
}

‎addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmCommunicationService.java‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import java.net.URL;
3333
import java.nio.file.Files;
3434
import java.nio.file.Paths;
35+
import java.util.HashMap;
3536
import java.util.List;
3637
import java.util.Map;
3738
import java.util.stream.Collectors;
@@ -75,6 +76,11 @@ public LlmCommunicationService(LlmOptions options) {
7576
requestor = new Requestor(HttpSender.MANUAL_REQUEST_INITIATOR, new HistoryPersister());
7677
}
7778

79+
/** For testing purposes only. */
80+
LlmCommunicationService(LlmAssistant assistant) {
81+
this.llmAssistant = assistant;
82+
}
83+
7884
private ChatLanguageModel buildModel(LlmOptions options) {
7985
return switch (options.getModelProvider()) {
8086
case AZURE_OPENAI ->
@@ -170,7 +176,13 @@ public void reviewAlert(Alert alert) {
170176
LOGGER.debug("Reviewing alert : {}", alert.getName());
171177
LOGGER.debug("Confidence level from ZAP : {}", alert.getConfidence());
172178
Stats.incCounter("stats.llm.alertreview.call");
173-
llmConfidence = llmAssistant.review(alert.getDescription(), alert.getEvidence());
179+
if (alert.getOtherInfo().isBlank()) {
180+
llmConfidence = llmAssistant.review(alert.getDescription(), alert.getEvidence());
181+
} else {
182+
llmConfidence =
183+
llmAssistant.review(
184+
alert.getDescription(), alert.getEvidence(), alert.getOtherInfo());
185+
}
174186

175187
if (llmConfidence.getLevel() == alert.getConfidence()) {
176188
Stats.incCounter("stats.llm.alertreview.result.same");
@@ -184,7 +196,7 @@ public void reviewAlert(Alert alert) {
184196
llmConfidence.getExplanation());
185197
updatedAlert.setConfidence(llmConfidence.getLevel());
186198
updatedAlert.setOtherInfo(getUpdatedOtherInfo(alert, llmConfidence));
187-
Map<String, String> alertTags = alert.getTags();
199+
Map<String, String> alertTags = new HashMap<>(alert.getTags());
188200

189201
alertTags.putIfAbsent(AI_REVIEWED_TAG_KEY, "");
190202
updatedAlert.setTags(alertTags);

‎addOns/llm/src/test/java/org/zaproxy/addon/llm/services/LlmCommunicationServiceUnitTest.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/*
2+
* Zed Attack Proxy (ZAP) and its related class files.
3+
*
4+
* ZAP is an HTTP/HTTPS proxy for assessing web application security.
5+
*
6+
* Copyright 2025 The ZAP Development Team
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this file except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
package org.zaproxy.addon.llm.services;
21+
22+
import static org.mockito.ArgumentMatchers.anyString;
23+
import static org.mockito.BDDMockito.given;
24+
import static org.mockito.Mockito.mock;
25+
import static org.mockito.Mockito.verify;
26+
27+
import org.junit.jupiter.api.BeforeAll;
28+
import org.junit.jupiter.api.Disabled;
29+
import org.junit.jupiter.api.Test;
30+
import org.junit.jupiter.params.ParameterizedTest;
31+
import org.junit.jupiter.params.provider.EmptySource;
32+
import org.junit.jupiter.params.provider.ValueSource;
33+
import org.parosproxy.paros.Constant;
34+
import org.parosproxy.paros.core.scanner.Alert;
35+
import org.zaproxy.addon.llm.communication.Confidence;
36+
import org.zaproxy.zap.testutils.TestUtils;
37+
import org.zaproxy.zap.utils.I18N;
38+
39+
/** Unit test for {@link LlmCommunicationService}. */
40+
class LlmCommunicationServiceUnitTest extends TestUtils {
41+
42+
private static final Confidence CONFIDENCE =
43+
new Confidence(Alert.CONFIDENCE_MEDIUM, "explanation");
44+
45+
@BeforeAll
46+
static void beforeAll() {
47+
Constant.messages = mock(I18N.class);
48+
}
49+
50+
@Disabled("Pending zaproxy/zap-extensions#6651")
51+
@ParameterizedTest
52+
@EmptySource
53+
@ValueSource(strings = {" ", "\t", "\r", "\n"})
54+
void shouldUseTwoParamReviewMethodWhenNoOtherInfo(String otherInfo) {
55+
// Given
56+
LlmAssistant assistant = mock();
57+
LlmCommunicationService service = new LlmCommunicationService(assistant);
58+
59+
Alert alert = createBaseAlert();
60+
alert.setOtherInfo(otherInfo);
61+
62+
given(assistant.review(anyString(), anyString())).willReturn(CONFIDENCE);
63+
// When
64+
service.reviewAlert(alert);
65+
// Then
66+
verify(assistant).review(anyString(), anyString());
67+
}
68+
69+
@Disabled("Pending zaproxy/zap-extensions#6651")
70+
@Test
71+
void shouldUseThreeParamReviewMethodWhenHasOtherInfo() {
72+
// Given
73+
LlmAssistant assistant = mock();
74+
LlmCommunicationService service = new LlmCommunicationService(assistant);
75+
76+
Alert alert = createBaseAlert();
77+
alert.setOtherInfo("other info");
78+
79+
given(assistant.review(anyString(), anyString(), anyString())).willReturn(CONFIDENCE);
80+
// When
81+
service.reviewAlert(alert);
82+
// Then
83+
verify(assistant).review(anyString(), anyString(), anyString());
84+
}
85+
86+
private static Alert createBaseAlert() {
87+
return Alert.builder()
88+
.setDescription("desc")
89+
.setEvidence("evidence")
90+
.setConfidence(Alert.CONFIDENCE_MEDIUM)
91+
.build();
92+
}
93+
}

0 commit comments

Comments
 (0)