Skip to content
Snippets Groups Projects
Commit 1244e4b5 authored by Robert May's avatar Robert May Committed by Jessie Young
Browse files

Add support for TOFA AI API

Changelog: added
parent 3c29b511
No related branches found
No related tags found
2 merge requests!122597doc/gitaly: Remove references to removed metrics,!118754Add support for TOFA AI API
......@@ -715,6 +715,18 @@ def self.kroki_formats_attributes
attr_encrypted :product_analytics_clickhouse_connection_string, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :product_analytics_configurator_connection_string, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :openai_api_key, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
# TOFA API integration settngs
attr_encrypted :tofa_client_library_args, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_client_library_class, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_client_library_create_credentials_method, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_client_library_fetch_access_token_method, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_credentials, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_host, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_request_json_keys, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_request_payload, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_response_json_keys, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_url, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
attr_encrypted :tofa_access_token_expires_in, encryption_options_base_32_aes_256_gcm.merge(encode: false, encode_iv: false)
validates :disable_feed_token,
inclusion: { in: [true, false], message: N_('must be a boolean value') }
......
---
name: tofa_experimentation
introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/118754
rollout_issue_url:
milestone: '16.0'
type: development
group: group::threat insights
default_enabled: false
# frozen_string_literal: true
class AddTofaApplicationSettings < Gitlab::Database::Migration[2.1]
def change
change_table(:application_settings, bulk: true) do |t|
t.column :encrypted_tofa_credentials, :binary
t.column :encrypted_tofa_credentials_iv, :binary
t.column :encrypted_tofa_host, :binary
t.column :encrypted_tofa_host_iv, :binary
t.column :encrypted_tofa_url, :binary
t.column :encrypted_tofa_url_iv, :binary
t.column :encrypted_tofa_response_json_keys, :binary
t.column :encrypted_tofa_response_json_keys_iv, :binary
t.column :encrypted_tofa_request_json_keys, :binary
t.column :encrypted_tofa_request_json_keys_iv, :binary
t.column :encrypted_tofa_request_payload, :binary
t.column :encrypted_tofa_request_payload_iv, :binary
t.column :encrypted_tofa_client_library_class, :binary
t.column :encrypted_tofa_client_library_class_iv, :binary
t.column :encrypted_tofa_client_library_args, :binary
t.column :encrypted_tofa_client_library_args_iv, :binary
t.column :encrypted_tofa_client_library_create_credentials_method, :binary
t.column :encrypted_tofa_client_library_create_credentials_method_iv, :binary
t.column :encrypted_tofa_client_library_fetch_access_token_method, :binary
t.column :encrypted_tofa_client_library_fetch_access_token_method_iv, :binary
t.column :encrypted_tofa_access_token_expires_in, :binary
t.column :encrypted_tofa_access_token_expires_in_iv, :binary
end
end
end
c400976f894b3d451bbf1a58e57f376cd86916680bb10623217b851a593cd4ea
\ No newline at end of file
......@@ -11805,6 +11805,28 @@ CREATE TABLE application_settings (
silent_mode_enabled boolean DEFAULT false NOT NULL,
package_metadata_purl_types smallint[] DEFAULT '{}'::smallint[],
ci_max_includes integer DEFAULT 150 NOT NULL,
encrypted_tofa_credentials bytea,
encrypted_tofa_credentials_iv bytea,
encrypted_tofa_host bytea,
encrypted_tofa_host_iv bytea,
encrypted_tofa_url bytea,
encrypted_tofa_url_iv bytea,
encrypted_tofa_response_json_keys bytea,
encrypted_tofa_response_json_keys_iv bytea,
encrypted_tofa_request_json_keys bytea,
encrypted_tofa_request_json_keys_iv bytea,
encrypted_tofa_request_payload bytea,
encrypted_tofa_request_payload_iv bytea,
encrypted_tofa_client_library_class bytea,
encrypted_tofa_client_library_class_iv bytea,
encrypted_tofa_client_library_args bytea,
encrypted_tofa_client_library_args_iv bytea,
encrypted_tofa_client_library_create_credentials_method bytea,
encrypted_tofa_client_library_create_credentials_method_iv bytea,
encrypted_tofa_client_library_fetch_access_token_method bytea,
encrypted_tofa_client_library_fetch_access_token_method_iv bytea,
encrypted_tofa_access_token_expires_in bytea,
encrypted_tofa_access_token_expires_in_iv bytea,
CONSTRAINT app_settings_container_reg_cleanup_tags_max_list_size_positive CHECK ((container_registry_cleanup_tags_service_max_list_size >= 0)),
CONSTRAINT app_settings_container_registry_pre_import_tags_rate_positive CHECK ((container_registry_pre_import_tags_rate >= (0)::numeric)),
CONSTRAINT app_settings_dep_proxy_ttl_policies_worker_capacity_positive CHECK ((dependency_proxy_ttl_group_policy_worker_capacity >= 0)),
......@@ -5,7 +5,8 @@ module Llm
module OpenAi
module Completions
class ExplainVulnerability
DEFAULT_ERROR = '{"error": "An unexpected error has occurred."}'
DEFAULT_ERROR = 'An unexpected error has occurred.'
DIGIT_REGEX = '/^\d$/'
def initialize(template_class)
@template_class = template_class
......@@ -15,14 +16,30 @@ def execute(user, vulnerability, _options)
template = template_class.new(vulnerability)
response = response_for(user, template)
::Gitlab::Llm::OpenAi::ResponseService
.new(user, vulnerability, response, options: {})
.execute(Gitlab::Llm::OpenAi::ResponseModifiers::Chat.new)
if Feature.enabled?(:tofa_experimentation, vulnerability.project)
json = Gitlab::Json.parse(response, symbolize_names: true)
GraphqlTriggers.ai_completion_response(
user.to_global_id,
vulnerability.to_global_id,
{
id: SecureRandom.uuid,
model_name: vulnerability.class.name,
response_body: json.dig(*tofa_json_args),
errors: tofa_errors(json)
}
)
else
::Gitlab::Llm::OpenAi::ResponseService
.new(user, vulnerability, response, options: {})
.execute(Gitlab::Llm::OpenAi::ResponseModifiers::Chat.new)
end
# or refresh token grant
rescue StandardError => error
Gitlab::ErrorTracking.track_exception(error)
::Gitlab::Llm::OpenAi::ResponseService
.new(user, vulnerability, DEFAULT_ERROR, options: {})
.new(user, vulnerability, { error: { message: DEFAULT_ERROR } }.to_json, options: {})
.execute
end
......@@ -31,9 +48,28 @@ def execute(user, vulnerability, _options)
attr_reader :template_class
def response_for(user, template)
Gitlab::Llm::OpenAi::Client
client_class
.new(user)
.chat(content: template.to_prompt, **template.options)
.chat(content: template.to_prompt, **template.options(client_class))
end
def client_class
if Feature.enabled?(:tofa_experimentation)
::Gitlab::Llm::Tofa::Client
else
::Gitlab::Llm::OpenAi::Client
end
end
def tofa_json_args
str = Gitlab::CurrentSettings.current_application_settings.tofa_response_json_keys
str.split.map do |s|
s.match?(DIGIT_REGEX) ? s.to_i : s.to_sym
end
end
def tofa_errors(json)
[json.dig(:error, :message)].compact
end
end
end
......
......@@ -13,10 +13,10 @@ def initialize(vulnerability)
@vulnerability = vulnerability
end
def options
{
max_tokens: MAX_TOKENS
}
def options(client)
return {} unless client == ::Gitlab::Llm::OpenAi::Client
{ max_tokens: MAX_TOKENS }
end
def to_prompt
......
# frozen_string_literal: true
module Gitlab
module Llm
module Tofa
class Client
include ::Gitlab::Llm::OpenAi::ExponentialBackoff
def initialize(_user, configuration = Configuration.new)
@configuration = configuration
end
def chat(content:, **options)
HTTParty.post( # rubocop: disable Gitlab/HTTParty
url,
headers: headers,
body: default_payload_for(content).merge(options).to_json
)
end
private
retry_methods_with_exponential_backoff :chat
attr_reader :configuration
delegate(
:access_token,
:host,
:tofa_request_json_keys,
:tofa_request_payload,
:url,
to: :configuration
)
def headers
{
"Accept" => "application/json",
"Authorization" => "Bearer #{access_token}",
"Host" => host,
"Content-Type" => "application/json"
}
end
def default_payload_for(content)
json = JSON.parse(tofa_request_payload) # rubocop: disable Gitlab/Json
json_keys = tofa_request_json_keys.split(' ')
json[json_keys[0]][0][json_keys[1]][0][json_keys[2]] = content
json
end
end
end
end
end
# frozen_string_literal: true
module Gitlab
module Llm
module Tofa
class Configuration
def access_token
Rails.cache.fetch(
:tofa_access_token,
expires_in: tofa_access_token_expires_in.to_i.seconds,
skip_nil: true
) do
fresh_token
end
end
def host
tofa_host
end
def url
tofa_url
end
private
def settings
@settings ||= Gitlab::CurrentSettings.current_application_settings
end
delegate(
:tofa_access_token_expires_in,
:tofa_client_library_args,
:tofa_client_library_class,
:tofa_client_library_create_credentials_method,
:tofa_client_library_fetch_access_token_method,
:tofa_credentials,
:tofa_host,
:tofa_request_json_keys,
:tofa_request_payload,
:tofa_url,
to: :settings
)
def fresh_token
client_library_class = tofa_client_library_class.constantize
client_library_args = string_to_hash(tofa_client_library_args)
first_key = client_library_args.first[0]
client_library_args[first_key] = StringIO.new(tofa_credentials)
create_credentials_method = tofa_client_library_create_credentials_method.to_sym
fetch_access_token_method = tofa_client_library_fetch_access_token_method.to_sym
# rubocop:disable GitlabSecurity/PublicSend
response = client_library_class.public_send(
create_credentials_method,
**client_library_args
).public_send(fetch_access_token_method)
# rubocop:enable GitlabSecurity/PublicSend
response["access_token"]
end
def string_to_hash(str)
hash = {}
str.split(", ").each do |pair|
key, value = pair.split(" ")
value = value == "nil" ? nil : value
hash[key.to_sym] = value
end
hash
end
end
end
end
end
......@@ -19,6 +19,8 @@
allow(GraphqlTriggers).to receive(:ai_completion_response)
vulnerability.finding.location['file'] = 'main.c'
vulnerability.finding.location['start_line'] = 1
stub_feature_flags(tofa_experimentation: false)
end
describe '#execute' do
......@@ -128,7 +130,7 @@
id: anything,
model_name: vulnerability.class.name,
response_body: '',
errors: ['An unexpected error has occurred.']
errors: [{ message: 'An unexpected error has occurred.' }]
}))
end
end
......
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