Skip to content

Add AIGW service class for resolve_vulnerability

Alejandro Rodríguez requested to merge resolve-vuln-aigw into master

What does this MR do and why?

Part of #482070, this MR adds the Gitlab::Llm::AiGateway::Completions::ResolveVulnerability completion which moves the prompt to the AI Gateway (see feat: add prompt for resolve_vulnerability (gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!1334 - merged)).

The changes in this MR are activated by the ai_gateway_agents feature flag.

This MR also extracts the shared code used by both Gitlab::Llm::AiGateway::Completions::ResolveVulnerability and Gitlab::Llm::Completions::ResolveVulnerability into a helper module called Gitlab::Llm::Completions::ResolveVulnerability::Helpers.

MR acceptance checklist

Please evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Screenshots or screen recordings

Screenshots are required for UI changes, and strongly recommended for all other merge requests.

Before After

How to set up and validate locally

  1. Setup - If needed, follow the instructions in this MR to set up Resolve Vulnerability with AI in the GDK - gitlab-com/content-sites/handbook!9430 (merged)
  2. First, ensure the original completion works by ensuring the ai_gateway_agents feature flag is disabled and testing Resolve Vulnerability for any supported CWE listed a project's vulnerability report.
  3. Tail the AI Gateway logs with gdk tail gitlab-ai-gateway.
  4. The path of the request in the logs should be path=/v1/proxy/anthropic/v1/messages if the ai_gateway_agents feature flag is disabled.
  5. Next, enable the ai_gateway_agents feature flag in the GDK.
  6. In the gitlab-ai-gateway project in the GDK change to the resolve-vulnerability-prompt branch (feat: add prompt for resolve_vulnerability (gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!1334 - merged)) and restart the AI gateway.
  7. Test Resolve Vulnerability for the same CWE you tested previously.
  8. Tail the AI Gateway logs with gdk tail gitlab-ai-gateway.
  9. In the logs, you should see Initializing prompt registry from local yaml
  10. The path of the request in the logs should be path=/v1/prompts/resolve_vulnerability if the ai_gateway_agents feature flag is disabled.
Edited by Darby Frey

Merge request reports