Merge branch '464043-instrumenting-ai-features-so-that-we-can-filter-by-client' into 'master'

Instrumenting AI features so that we can filter by client

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/170643



Merged-by: Mark Chao <mchao@gitlab.com>
Approved-by: Jeff Park <jpark@gitlab.com>
Approved-by: Jonas Larsen <jlarsen@gitlab.com>
Approved-by: Mark Chao <mchao@gitlab.com>
Reviewed-by: Mark Chao <mchao@gitlab.com>
Reviewed-by: Jonas Larsen <jlarsen@gitlab.com>
Reviewed-by: Duo Code Reviewer <duo-code-review-bot@gitlab.com>
Co-authored-by: lesley <lrazzaghian@gitlab.com>

ee/app/services/llm/execute_method_service.rb

@@ -26,6 +26,8 @@ def execute
     def track_snowplow_event(result)
       request_id = result[:ai_message]&.request_id
 
+      client = Gitlab::Llm::Tracking.client_for_user_agent(options[:user_agent])
+
       Gitlab::Tracking.event(
         self.class.to_s,
         "execute_llm_method",
@@ -34,7 +36,8 @@ def track_snowplow_event(result)
         user: user,
         namespace: namespace,
         project: project,
-        requestId: request_id
+        requestId: request_id,
+        client: client
       )
     end
   end

ee/config/events/execute_llm_method.yml

+---
+description: Tracks when a call was made to an LLM
+internal_events: true
+action: execute_llm_method
+identifiers:
+- project
+- namespace
+- user
+additional_properties:
+  label:
+    description: The type of LLM request (such as 'chat' or 'categorize_question')
+  property:
+    description: Success of request
+  client:
+    description: Client such as 'web' or 'vscode'
+  requestId:
+    description: Request id
+product_group: duo_chat
+milestone: '16.0'
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/117877
+distributions:
+- ee
+tiers:
+- free
+- premium
+- ultimate

ee/lib/gitlab/llm/tracking.rb

@@ -3,6 +3,21 @@
 module Gitlab
   module Llm
     class Tracking
+      USER_AGENT_CLIENTS = {
+        UsageDataCounters::VSCodeExtensionActivityUniqueCounter::VS_CODE_USER_AGENT_REGEX =>
+          'vscode',
+        UsageDataCounters::JetBrainsPluginActivityUniqueCounter::JETBRAINS_USER_AGENT_REGEX =>
+          'jetbrains',
+        UsageDataCounters::JetBrainsBundledPluginActivityUniqueCounter::JETBRAINS_BUNDLED_USER_AGENT_REGEX =>
+          'jetbrains_bundled',
+        UsageDataCounters::VisualStudioExtensionActivityUniqueCounter::VISUAL_STUDIO_EXTENSION_USER_AGENT_REGEX =>
+          'visual_studio',
+        UsageDataCounters::NeovimPluginActivityUniqueCounter::NEOVIM_PLUGIN_USER_AGENT_REGEX =>
+          'neovim',
+        UsageDataCounters::GitLabCliActivityUniqueCounter::GITLAB_CLI_USER_AGENT_REGEX =>
+          'gitlab_cli'
+      }.freeze
+
       def self.event_for_ai_message(category, action, ai_message:)
         ::Gitlab::Tracking.event(
           category,
@@ -17,9 +32,8 @@ def self.event_for_ai_message(category, action, ai_message:)
       def self.client_for_user_agent(user_agent)
         return unless user_agent.present?
 
-        user_agent.match?(Gitlab::Regex.vs_code_user_agent_regex) ? 'vscode' : 'web'
+        USER_AGENT_CLIENTS.find { |regex, _client| user_agent.match?(regex) }&.last || 'web'
       end
-      private_class_method :client_for_user_agent
     end
   end
 end

ee/spec/lib/gitlab/llm/tracking_spec.rb

@@ -6,8 +6,8 @@
   let(:user) { build(:user) }
   let(:resource) { build(:project) }
   let(:ai_action_name) { 'chat' }
-  let(:user_agent) { nil }
   let(:request_id) { 'uuid' }
+  let(:user_agent) { nil }
 
   let(:ai_message) do
     build(:ai_message,
@@ -32,39 +32,67 @@
         client: nil
       )
     end
+  end
+
+  describe '.client_for_user_agent' do
+    subject { described_class.client_for_user_agent(user_agent) }
 
-    context 'with browser user agent' do
+    context 'when user agent is from a web browser' do
       let(:user_agent) { 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)' }
 
-      it 'tracks event with correct params' do
-        event_for_ai_message
-
-        expect_snowplow_event(
-          category: 'Category',
-          action: 'my_action',
-          label: ai_action_name,
-          property: request_id,
-          user: user,
-          client: 'web'
-        )
-      end
+      it { is_expected.to eq('web') }
     end
 
-    context 'with vscode user agent' do
+    context 'when user agent is from VS Code' do
       let(:user_agent) { 'vs-code-gitlab-workflow/3.11.1 VSCode/1.52.1 Node.js/12.14.1 (darwin; x64)' }
 
-      it 'tracks event with correct params' do
-        event_for_ai_message
-
-        expect_snowplow_event(
-          category: 'Category',
-          action: 'my_action',
-          label: ai_action_name,
-          property: request_id,
-          user: user,
-          client: 'vscode'
-        )
+      it { is_expected.to eq('vscode') }
+    end
+
+    context 'when user agent is from JetBrains plugin' do
+      let(:user_agent) { 'gitlab-jetbrains-plugin/0.0.1 intellij-idea/2021.2.4 java/11.0.13 mac-os-x/aarch64/12.1' }
+
+      it { is_expected.to eq('jetbrains') }
+    end
+
+    context 'when user agent is from JetBrains bundled plugin' do
+      let(:user_agent) do
+        'IntelliJ-GitLab-Plugin PhpStorm/PS-232.6734.11 (JRE 17.0.7+7-b966.2; Linux 6.2.0-20-generic; amd64)'
       end
+
+      it { is_expected.to eq('jetbrains_bundled') }
+    end
+
+    context 'when user agent is from Visual Studio extension' do
+      let(:user_agent) { 'code-completions-language-server-experiment (gl-visual-studio-extension:1.0.0.0; arch:X64;)' }
+
+      it { is_expected.to eq('visual_studio') }
+    end
+
+    context 'when user agent is from Neovim plugin' do
+      let(:user_agent) do
+        'code-completions-language-server-experiment (Neovim:0.9.0; gitlab.vim (v0.1.0); arch:amd64; os:darwin)'
+      end
+
+      it { is_expected.to eq('neovim') }
+    end
+
+    context 'when user agent is from GitLab CLI (old format)' do
+      let(:user_agent) { 'GLab - GitLab CLI' }
+
+      it { is_expected.to eq('gitlab_cli') }
+    end
+
+    context 'when user agent is from GitLab CLI (current format)' do
+      let(:user_agent) { 'glab/v1.25.3-27-g7ec258fb (built 2023-02-16), darwin' }
+
+      it { is_expected.to eq('gitlab_cli') }
+    end
+
+    context 'when user agent is nil' do
+      let(:user_agent) { nil }
+
+      it { is_expected.to be_nil }
     end
   end
 end

ee/spec/services/llm/execute_method_service_spec.rb

@@ -13,6 +13,7 @@
   let(:success) { true }
   let(:request_id) { 'uuid' }
   let(:chat_message) { instance_double(Gitlab::Llm::ChatMessage, request_id: request_id) }
+  let(:client) { nil }
   let(:service_response) do
     instance_double(
       ServiceResponse,
@@ -94,7 +95,8 @@
           user: user,
           namespace: group,
           project: project,
-          requestId: 'uuid'
+          requestId: 'uuid',
+          client: nil
         }
       end
 
@@ -102,6 +104,9 @@
         allow_next_instance_of(service_class, user, resource, options) do |instance|
           allow(instance).to receive(:execute).and_return(service_response)
         end
+
+        allow(Gitlab::Llm::Tracking).to receive(:client_for_user_agent)
+                                          .and_return(client)
       end
 
       shared_examples 'successful tracking' do
@@ -166,6 +171,13 @@
 
         it_behaves_like 'successful tracking'
       end
+
+      context 'when client is detected' do
+        let(:client) { 'web' }
+        let(:expected_params) { default_params.merge(client: 'web') }
+
+        it_behaves_like 'successful tracking'
+      end
     end
   end
 end

lib/gitlab/regex.rb

@@ -282,10 +282,6 @@ def ml_model_version_name_regex
     def ml_model_file_name_regex
       @ml_model_file_name_regex ||= %r{\A[A-Za-z0-9\.\_\-\+ ]+\z}
     end
-
-    def vs_code_user_agent_regex
-      /\Avs-code-gitlab-workflow/
-    end
   end
 end
 

lib/gitlab/usage_data_counters/vscode_extension_activity_unique_counter.rb

@@ -4,11 +4,11 @@ module Gitlab
   module UsageDataCounters
     module VSCodeExtensionActivityUniqueCounter
       VS_CODE_API_REQUEST_ACTION = 'i_code_review_user_vs_code_api_request'
+      VS_CODE_USER_AGENT_REGEX = /\Avs-code-gitlab-workflow/
 
       class << self
         def track_api_request_when_trackable(user_agent:, user:)
-          user_agent&.match?(Gitlab::Regex.vs_code_user_agent_regex) && track_unique_action_by_user(
-            VS_CODE_API_REQUEST_ACTION, user)
+          user_agent&.match?(VS_CODE_USER_AGENT_REGEX) && track_unique_action_by_user(VS_CODE_API_REQUEST_ACTION, user)
         end
 
         private