From 72c2ba0942313a8e56c41fedd8000d8374097878 Mon Sep 17 00:00:00 2001 From: Najam Ul Saqib Date: Fri, 18 Jul 2025 17:26:03 +0500 Subject: [PATCH] pass otherinfo to llm Signed-off-by: Najam Ul Saqib --- .../addon/llm/services/LlmAssistant.java | 57 ++++++++++++----- .../llm/services/LlmCommunicationService.java | 16 ++++- .../LlmCommunicationServiceUnitTest.java | 62 +++++++++++++++++++ 3 files changed, 116 insertions(+), 19 deletions(-) diff --git a/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmAssistant.java b/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmAssistant.java index 5870bb3db7d..10c05cb1c3d 100644 --- a/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmAssistant.java +++ b/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmAssistant.java @@ -34,22 +34,45 @@ public interface LlmAssistant { "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") HttpRequestList complete(); - @SystemMessage( - "You are a web application security expert reviewing potential false positives. Answer only in JSON.") - @UserMessage( - "Your task is to review the following finding from ZAP (Zed Attack Proxy).\n" - + "The confidence level is a pull down field which allows you to specify how confident you are in the validity of the finding : \n" - + "- 0 if it's False Positive\n" - + "- 1 if it's Low\n" - + "- 2 if it's Medium\n" - + "- 3 if it's High\n" - + "\n" - + "The alert is described as follows : {{description}}\n" - + "\n" - + "As evidence, the HTTP response contains :\n" - + "---\n" - + "{{evidence}}\n" - + "---\n" - + "Provide a short consistent explanation of the new score.\n") + static final String ALERT_REVIEW_SYSTEM_MSG = + "You are a web application security expert reviewing potential false positives. Answer only in JSON."; + static final String ALERT_REVIEW_GOAL = + "Provide a short consistent explanation of the new score.\n"; + static final String ALERT_REVIEW_PROMPT = + """ + Your task is to review the following finding from ZAP (Zed Attack Proxy). + The confidence level is a pull down field which allows you to specify how confident you are in the validity of the finding: + - 0 if it's False Positive + - 1 if it's Low + - 2 if it's Medium + - 3 if it's High + The alert is described as follows: {{description}} + + As evidence, the HTTP message contains: + --- + {{evidence}} + --- + """ + + ALERT_REVIEW_GOAL; + + static final String ALERT_REVIEW_OTHERINFO_PROMPT = + ALERT_REVIEW_PROMPT + + """ + Also, here's some additional information that may be useful for you to reach your conclusion: + --- + {{otherinfo}} + --- + """ + + ALERT_REVIEW_GOAL; + + @SystemMessage(ALERT_REVIEW_SYSTEM_MSG) + @UserMessage(ALERT_REVIEW_PROMPT) Confidence review(@V("description") String description, @V("evidence") String evidence); + + @SystemMessage(ALERT_REVIEW_SYSTEM_MSG) + @UserMessage(ALERT_REVIEW_OTHERINFO_PROMPT) + Confidence review( + @V("description") String description, + @V("evidence") String evidence, + @V("otherinfo") String otherinfo); } diff --git a/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmCommunicationService.java b/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmCommunicationService.java index 9c41a10b358..e51ee29f071 100644 --- a/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmCommunicationService.java +++ b/addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmCommunicationService.java @@ -32,6 +32,7 @@ import java.net.URL; import java.nio.file.Files; import java.nio.file.Paths; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -75,6 +76,11 @@ public LlmCommunicationService(LlmOptions options) { requestor = new Requestor(HttpSender.MANUAL_REQUEST_INITIATOR, new HistoryPersister()); } + /** For testing purposes only. */ + LlmCommunicationService(LlmAssistant assistant) { + this.llmAssistant = assistant; + } + private ChatLanguageModel buildModel(LlmOptions options) { return switch (options.getModelProvider()) { case AZURE_OPENAI -> @@ -170,7 +176,13 @@ public void reviewAlert(Alert alert) { LOGGER.debug("Reviewing alert : {}", alert.getName()); LOGGER.debug("Confidence level from ZAP : {}", alert.getConfidence()); Stats.incCounter("stats.llm.alertreview.call"); - llmConfidence = llmAssistant.review(alert.getDescription(), alert.getEvidence()); + if (alert.getOtherInfo().isBlank()) { + llmConfidence = llmAssistant.review(alert.getDescription(), alert.getEvidence()); + } else { + llmConfidence = + llmAssistant.review( + alert.getDescription(), alert.getEvidence(), alert.getOtherInfo()); + } if (llmConfidence.getLevel() == alert.getConfidence()) { Stats.incCounter("stats.llm.alertreview.result.same"); @@ -184,7 +196,7 @@ public void reviewAlert(Alert alert) { llmConfidence.getExplanation()); updatedAlert.setConfidence(llmConfidence.getLevel()); updatedAlert.setOtherInfo(getUpdatedOtherInfo(alert, llmConfidence)); - Map alertTags = alert.getTags(); + Map alertTags = new HashMap<>(alert.getTags()); alertTags.putIfAbsent(AI_REVIEWED_TAG_KEY, ""); updatedAlert.setTags(alertTags); diff --git a/addOns/llm/src/test/java/org/zaproxy/addon/llm/services/LlmCommunicationServiceUnitTest.java b/addOns/llm/src/test/java/org/zaproxy/addon/llm/services/LlmCommunicationServiceUnitTest.java index c104726792f..3e8b67292c6 100644 --- a/addOns/llm/src/test/java/org/zaproxy/addon/llm/services/LlmCommunicationServiceUnitTest.java +++ b/addOns/llm/src/test/java/org/zaproxy/addon/llm/services/LlmCommunicationServiceUnitTest.java @@ -21,16 +21,35 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.hasEntry; import static org.hamcrest.Matchers.is; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.BDDMockito.given; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import java.util.Map; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.parosproxy.paros.Constant; import org.parosproxy.paros.core.scanner.Alert; +import org.zaproxy.addon.llm.communication.Confidence; import org.zaproxy.zap.testutils.TestUtils; +import org.zaproxy.zap.utils.I18N; /** Unit test for {@link LlmCommunicationService}. */ class LlmCommunicationServiceUnitTest extends TestUtils { + private static final Confidence CONFIDENCE = + new Confidence(Alert.CONFIDENCE_MEDIUM, "explanation"); + + @BeforeAll + static void beforeAll() { + Constant.messages = mock(I18N.class); + } + @Test void shouldNotBeConsideredReviewdIfNoTags() { // Given @@ -62,4 +81,47 @@ void shouldBeConsideredReviewdIfMarkedAsSuch() { // Then assertThat(result, is(equalTo(true))); } + + @ParameterizedTest + @ValueSource(strings = {"", " ", "\t", "\r", "\n"}) + void shouldUseTwoParamReviewMethodWhenNoOtherInfo(String otherInfo) { + // Given + LlmAssistant assistant = mock(); + LlmCommunicationService service = new LlmCommunicationService(assistant); + + Alert alert = createBaseAlert(); + alert.setOtherInfo(otherInfo); + + given(assistant.review(anyString(), anyString())).willReturn(CONFIDENCE); + // When + service.reviewAlert(alert); + // Then + verify(assistant).review(anyString(), anyString()); + assertThat(alert.getTags(), hasEntry(LlmCommunicationService.AI_REVIEWED_TAG_KEY, "")); + } + + @Test + void shouldUseThreeParamReviewMethodWhenHasOtherInfo() { + // Given + LlmAssistant assistant = mock(); + LlmCommunicationService service = new LlmCommunicationService(assistant); + + Alert alert = createBaseAlert(); + alert.setOtherInfo("other info"); + + given(assistant.review(anyString(), anyString(), anyString())).willReturn(CONFIDENCE); + // When + service.reviewAlert(alert); + // Then + verify(assistant).review(anyString(), anyString(), anyString()); + assertThat(alert.getTags(), hasEntry(LlmCommunicationService.AI_REVIEWED_TAG_KEY, "")); + } + + private static Alert createBaseAlert() { + return Alert.builder() + .setDescription("desc") + .setEvidence("evidence") + .setConfidence(Alert.CONFIDENCE_MEDIUM) + .build(); + } }