feat: Add support for disabling LLM monitoring at the account level. (#…

…2592)
newrelic · Jul 8, 2024 · c31451a · c31451a
1 parent b17bfab
commit c31451a
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 5 deletions.
diff --git a/src/Agent/NewRelic/Agent/Core/Configuration/DefaultConfiguration.cs b/src/Agent/NewRelic/Agent/Core/Configuration/DefaultConfiguration.cs
@@ -2076,8 +2076,8 @@ public bool AiMonitoringEnabled
         {
             get
             {
-                // AI Monitoring is disabled in High Security Mode
-                return !HighSecurityModeEnabled && EnvironmentOverrides(_localConfiguration.aiMonitoring.enabled, "NEW_RELIC_AI_MONITORING_ENABLED");
+                // AI Monitoring is disabled in High Security Mode and can be disabled at the account level
+                return !HighSecurityModeEnabled && ServerCanDisable(_serverConfiguration.AICollectionEnabled, EnvironmentOverrides(_localConfiguration.aiMonitoring.enabled, "NEW_RELIC_AI_MONITORING_ENABLED"));
             }
         }
 

diff --git a/src/Agent/NewRelic/Agent/Core/Configuration/ReportedConfiguration.cs b/src/Agent/NewRelic/Agent/Core/Configuration/ReportedConfiguration.cs
@@ -665,7 +665,9 @@ public ReportedConfiguration(IConfiguration configuration)
         [JsonProperty("agent.disable_file_system_watcher")]
         public bool DisableFileSystemWatcher => _configuration.DisableFileSystemWatcher;
 
-        [JsonProperty("agent.ai_monitoring.enabled")]
+        // To support account level disable feature, the name cannot include the "agent" prefix.
+        // The account level disable feature looks for this setting and will only reply with the ServerConfiguration setting if it is present.
+        [JsonProperty("ai_monitoring.enabled")]
         public bool AiMonitoringEnabled => _configuration.AiMonitoringEnabled;
 
         [JsonProperty("ai_monitoring.streaming.enabled")]

diff --git a/src/Agent/NewRelic/Agent/Core/Configuration/ServerConfiguration.cs b/src/Agent/NewRelic/Agent/Core/Configuration/ServerConfiguration.cs
@@ -42,6 +42,9 @@ public class ServerConfiguration
         [JsonProperty("collect_traces")]
         public bool? TraceCollectionEnabled { get; set; }
 
+        [JsonProperty("collect_ai")]
+        public bool? AICollectionEnabled { get; set; }
+
         [JsonProperty("data_report_period")]
         public long? DataReportPeriod { get; set; }
 

diff --git a/...egrationTests/IntegrationTestHelpers/RemoteServiceFixtures/ConsoleDynamicMethodFixture.cs b/...egrationTests/IntegrationTestHelpers/RemoteServiceFixtures/ConsoleDynamicMethodFixture.cs
@@ -16,6 +16,16 @@ public ConsoleDynamicMethodFixtureFWLatest()
         }
     }
 
+    /// <summary>
+    /// Use this fixture for AIM Account level disabled tests
+    /// </summary>
+    public class ConsoleDynamicMethodFixtureFWLatestAIM : ConsoleDynamicMethodFixtureFW481
+    {
+        public override string TestSettingCategory { get { return "AIM"; } }
+        public ConsoleDynamicMethodFixtureFWLatestAIM()
+        {
+        }
+    }
 
     /// <summary>
     /// Use this fixture for High Security Mode tests
@@ -121,6 +131,18 @@ public ConsoleDynamicMethodFixtureCoreLatest()
         }
     }
 
+    /// <summary>
+    /// Use this fixture for AIM Account level disabled tests
+    /// </summary>
+    public class ConsoleDynamicMethodFixtureCoreLatestAIM : ConsoleDynamicMethodFixtureCore80
+    {
+        public override string TestSettingCategory { get { return "AIM"; } }
+        public ConsoleDynamicMethodFixtureCoreLatestAIM()
+        {
+        }
+
+    }
+
     /// <summary>
     /// Use this fixture for High Security Mode tests
     /// </summary>

diff --git a/tests/Agent/IntegrationTests/IntegrationTests/LLM/LLMAccountDisabledTests.cs b/tests/Agent/IntegrationTests/IntegrationTests/LLM/LLMAccountDisabledTests.cs
@@ -0,0 +1,82 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using NewRelic.Agent.IntegrationTestHelpers;
+using NewRelic.Agent.IntegrationTestHelpers.RemoteServiceFixtures;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace NewRelic.Agent.IntegrationTests.LLM
+{
+    public abstract class LlmAccountDisabledTestsBase<TFixture> : NewRelicIntegrationTest<TFixture>
+where TFixture : ConsoleDynamicMethodFixture
+    {
+        private readonly TFixture _fixture;
+        private const string _model = "meta13";
+        private string _prompt = "In one sentence, what is a large-language model?";
+
+        public LlmAccountDisabledTestsBase(TFixture fixture, ITestOutputHelper output) : base(fixture)
+        {
+            _fixture = fixture;
+            _fixture.SetTimeout(TimeSpan.FromMinutes(2));
+            _fixture.TestLogger = output;
+            _fixture.AddActions(
+                setupConfiguration: () =>
+                {
+                    new NewRelicConfigModifier(fixture.DestinationNewRelicConfigFilePath)
+                        .ForceTransactionTraces()
+                        .EnableAiMonitoring(true) // must be true to test override.
+                        .SetLogLevel("finest");
+                },
+                exerciseApplication: () =>
+                {
+                    _fixture.AgentLog.WaitForLogLines(AgentLogBase.MetricDataLogLineRegex, TimeSpan.FromMinutes(3), 2);
+                }
+            );
+
+            _fixture.AddCommand($"LLMExerciser InvokeModel {_model} {LLMHelpers.ConvertToBase64(_prompt)}");
+
+            _fixture.Initialize();
+        }
+
+        [Fact]
+        public void BedrockDisabledTest()
+        {
+            // Make sure it actually got called
+            var transactionEvent = _fixture.AgentLog.TryGetTransactionEvent($"OtherTransaction/Custom/MultiFunctionApplicationHelpers.NetStandardLibraries.LLM.LLMExerciser/InvokeModel");
+            Assert.NotNull(transactionEvent);
+
+            var unexpectedMetrics = new List<Assertions.ExpectedMetric>
+            {
+                new Assertions.ExpectedMetric { metricName = @"Supportability/DotNet/ML/Bedrock/.*", IsRegexName = true },
+                new Assertions.ExpectedMetric { metricName = @"Custom/Llm/.*", IsRegexName = true },
+            };
+
+            var metrics = _fixture.AgentLog.GetMetrics().ToList();
+            Assertions.MetricsDoNotExist(unexpectedMetrics, metrics);
+
+
+        }
+    }
+    [NetCoreTest]
+    public class LlmAccountDisabledTest_CoreLatest : LlmAccountDisabledTestsBase<ConsoleDynamicMethodFixtureCoreLatestAIM>
+    {
+        public LlmAccountDisabledTest_CoreLatest(ConsoleDynamicMethodFixtureCoreLatestAIM fixture, ITestOutputHelper output)
+            : base(fixture, output)
+        {
+        }
+    }
+
+    [NetFrameworkTest]
+    public class LlmAccountDisabledTest_FWLatest : LlmAccountDisabledTestsBase<ConsoleDynamicMethodFixtureFWLatestAIM>
+    {
+        public LlmAccountDisabledTest_FWLatest(ConsoleDynamicMethodFixtureFWLatestAIM fixture, ITestOutputHelper output)
+            : base(fixture, output)
+        {
+        }
+    }
+
+}
diff --git a/tests/Agent/UnitTests/Core.UnitTest/DataTransport/AgentSettingsTests.cs b/tests/Agent/UnitTests/Core.UnitTest/DataTransport/AgentSettingsTests.cs
@@ -338,7 +338,7 @@ public void serializes_correctly()
                         }
                     ],
                     "agent.disable_file_system_watcher": false,
-                    "agent.ai_monitoring.enabled": true,
+                    "ai_monitoring.enabled": true,
                     "ai_monitoring.streaming.enabled": true,
                     "ai_monitoring.record_content.enabled": true
                 }

diff --git a/tests/Agent/UnitTests/Core.UnitTest/DataTransport/ConnectModelTests.cs b/tests/Agent/UnitTests/Core.UnitTest/DataTransport/ConnectModelTests.cs
@@ -408,7 +408,7 @@ public void serializes_correctly()
                                 }
                             ],
                             "agent.disable_file_system_watcher": false,
-                            "agent.ai_monitoring.enabled": true,
+                            "ai_monitoring.enabled": true,
                             "ai_monitoring.streaming.enabled": true,
                             "ai_monitoring.record_content.enabled": true
                         },