Skip to content
Snippets Groups Projects
Unverified Commit be00b33c authored by Gregory Havenga's avatar Gregory Havenga Committed by Mehmet Emin Inaç
Browse files

Ensure `include_source_code` is appropriately passed to the template

Additionally remove the method default to prevent this issue occurring
again.

Changelog: fixed
EE: true
parent e22de923
No related branches found
No related tags found
No related merge requests found
Loading
Loading
@@ -46,7 +46,7 @@ def formatted_error_response(message)
 
def response_for(user, vulnerability, options)
Rails.cache.fetch(cache_key(vulnerability, options), expires_in: 5.minutes, skip_nil: true) do
prompt = ai_prompt_class.new(vulnerability, options).to_prompt
prompt = ai_prompt_class.new(vulnerability).to_prompt(include_source_code: options[:include_source_code])
if prompt.to_s.empty?
formatted_error_response(NULL_PROMPT_ERROR)
else
Loading
Loading
Loading
Loading
@@ -34,8 +34,19 @@ def initialize(vulnerability, params = {})
@params = params
end
 
def to_prompt(include_source_code: nil)
prompt(include_source_code: include_source_code)
# This will always provide a prompt by default, but if
# `include_source_code` is set true, we'll respond with nil if no
# code is available to provide a better interface to the
# frontend.
def to_prompt(include_source_code:)
case include_source_code
when true
prompt_with_code if eligible_code?
when false
prompt_without_code
else
eligible_and_safe_code? ? prompt_with_code : prompt_without_code
end
end
 
def presubmission_checks
Loading
Loading
@@ -52,21 +63,6 @@ def presubmission_checks
delegate :title, :description, :file, to: :vulnerability
delegate :source_code?, :vulnerable_code, to: :finding
 
# This will always provide a prompt by default, but if
# `include_source_c45ode` is set, we'll respond with nil if no
# code is available to provide a better interface to the
# frontend.
def prompt(include_source_code: false)
case include_source_code
when true
prompt_with_code if eligible_code?
when false
prompt_without_code
else
eligible_and_safe_code? ? prompt_with_code : prompt_without_code
end
end
# rubocop: disable CodeReuse/ActiveRecord
def identifiers
finding.identifiers.pluck(:name).join(", ")
Loading
Loading
Loading
Loading
@@ -3,6 +3,8 @@
require 'spec_helper'
 
RSpec.describe Gitlab::Llm::Completions::ExplainVulnerability, feature_category: :vulnerability_management do
using RSpec::Parameterized::TableSyntax
let(:prompt_class) { Gitlab::Llm::Templates::ExplainVulnerability }
let(:example_answer) { "Sure, ..." }
let(:example_response) do
Loading
Loading
@@ -66,12 +68,12 @@
context 'when a null prompt is returned by the template class' do
before do
allow_next_instance_of(prompt_class) do |prompt_class|
allow(prompt_class).to receive(:to_prompt).and_return(nil)
allow(prompt_class).to receive(:to_prompt).with(include_source_code: true).and_return(nil)
end
end
 
it 'returns the default error response' do
explain.execute(user, vulnerability, {})
explain.execute(user, vulnerability, { include_source_code: true })
 
expect(GraphqlTriggers).to have_received(:ai_completion_response)
.with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
Loading
Loading
@@ -226,6 +228,34 @@
end
end
 
context 'when the `include_source_code` option is passed' do
where(:include_source_code) do
[
true,
false,
nil
]
end
with_them do
before do
allow_next_instance_of(llm_client) do |client|
allow(client).to receive(client_method).once.and_return(example_response[llm_client].to_json)
end
end
it 'passes the parameter to the template class appropriately' do
expect_next_instance_of(prompt_class) do |instance|
expect(instance).to receive(:to_prompt)
.with(include_source_code: include_source_code)
.and_call_original
end
explain.execute(user, vulnerability, { include_source_code: include_source_code })
end
end
end
context 'when the `include_source_code` option is toggled' do
before do
allow_next_instance_of(llm_client) do |client|
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment