summary: AWS Bedrock instrumentation

.github/workflows/all_solutions.yml

@@ -247,7 +247,8 @@ jobs:
           DistributedTracing, 
           Errors, 
           HttpClientInstrumentation, 
-          InfiniteTracing, 
+          InfiniteTracing,
+          LLM,
           Logging.ContextData, 
           Logging.HsmAndCsp, 
           Logging.LocalDecoration, 

.github/workflows/run_integration_tests.yml

@@ -53,7 +53,7 @@ jobs:
           if [ "${{ inputs.integration-test-namespaces }}" == "ALL" ] ; then
             # Use the full list of namespaces
             namespaces="[ 'AgentFeatures', 'AgentLogs', 'AgentMetrics', 'Api', 'AppDomainCaching', 'AspNetCore', 'BasicInstrumentation', 'CatInbound', 'CatOutbound', 'CodeLevelMetrics', 'Configuration', \
-          'CSP', 'CustomAttributes', 'CustomInstrumentation', 'DataTransmission', 'DistributedTracing', 'Errors', 'HttpClientInstrumentation', 'InfiniteTracing', 'Logging.ContextData', \
+          'CSP', 'CustomAttributes', 'CustomInstrumentation', 'DataTransmission', 'DistributedTracing', 'Errors', 'HttpClientInstrumentation', 'InfiniteTracing', 'LLM', 'Logging.ContextData', \
           'Logging.HsmAndCsp', 'Logging.LocalDecoration', 'Logging.LogLevelDetection', 'Logging.MaxSamplesStored', 'Logging.MetricsAndForwarding', 'Logging.ZeroMaxSamplesStored', \
           'Owin', 'MassTransit', 'ReJit.NetCore', 'ReJit.NetFramework', 'RequestHandling', 'RequestHeadersCapture.AspNet', 'RequestHeadersCapture.AspNetCore', 'RequestHeadersCapture.EnvironmentVariables', \
           'RequestHeadersCapture.Owin', 'RequestHeadersCapture.WCF', 'RestSharp', 'WCF.Client.IIS.ASPDisabled', 'WCF.Client.IIS.ASPEnabled', 'WCF.Client.Self', \

FullAgent.sln

@@ -178,6 +178,7 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Home", "src\Agent\NewRelic\
 		{87EF7419-8390-49F2-8C51-4DF9B886E834} = {87EF7419-8390-49F2-8C51-4DF9B886E834}
 		{8D2B52DD-D45C-481D-92F0-28990F168820} = {8D2B52DD-D45C-481D-92F0-28990F168820}
 		{9784EAEA-C32F-4DC4-BD84-4075BBA541AB} = {9784EAEA-C32F-4DC4-BD84-4075BBA541AB}
+		{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A} = {9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A}
 		{9E9E1367-E09E-45F9-9A01-7D4489FBCB23} = {9E9E1367-E09E-45F9-9A01-7D4489FBCB23}
 		{A4B357E3-9CEF-492C-A0E2-2397D15F1CD6} = {A4B357E3-9CEF-492C-A0E2-2397D15F1CD6}
 		{AA683341-1FF5-4D45-A831-1BAF3C100A5C} = {AA683341-1FF5-4D45-A831-1BAF3C100A5C}
@@ -218,6 +219,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Kafka", "src\Agent\NewRelic
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "AspNetCore6Plus", "src\Agent\NewRelic\Agent\Extensions\Providers\Wrapper\AspNetCore6Plus\AspNetCore6Plus.csproj", "{D4F48A7F-F3D3-4303-921D-BF7FE34B7118}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Bedrock", "src\Agent\NewRelic\Agent\Extensions\Providers\Wrapper\Bedrock\Bedrock.csproj", "{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -456,6 +459,10 @@ Global
 		{D4F48A7F-F3D3-4303-921D-BF7FE34B7118}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{D4F48A7F-F3D3-4303-921D-BF7FE34B7118}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{D4F48A7F-F3D3-4303-921D-BF7FE34B7118}.Release|Any CPU.Build.0 = Release|Any CPU
+		{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -526,6 +533,7 @@ Global
 		{F921A316-A8D2-4992-B894-69A3B5CA34E3} = {5E86E10A-C38F-48CB-ADE9-67B22BB2F50A}
 		{270A9CC8-8031-49F4-A380-1389E7517DB7} = {5E86E10A-C38F-48CB-ADE9-67B22BB2F50A}
 		{D4F48A7F-F3D3-4303-921D-BF7FE34B7118} = {5E86E10A-C38F-48CB-ADE9-67B22BB2F50A}
+		{9C20BC4E-7A9F-4518-B3E7-C1FFD6C0EC8A} = {5E86E10A-C38F-48CB-ADE9-67B22BB2F50A}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		EnterpriseLibraryConfigurationToolBinariesPath = packages\Unity.2.1.505.2\lib\NET35

build/ArtifactBuilder/CoreAgentComponents.cs

@@ -59,6 +59,7 @@ protected override void CreateAgentComponents()
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.MassTransitLegacy.dll",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Kafka.dll",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.AspNetCore6Plus.dll",
+                $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Bedrock.dll",
             };
 
             var wrapperXmls = new[]
@@ -82,6 +83,7 @@ protected override void CreateAgentComponents()
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.MassTransitLegacy.Instrumentation.xml",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Kafka.Instrumentation.xml",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.AspNetCore6Plus.Instrumentation.xml",
+                $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Bedrock.Instrumentation.xml",
             };
 
             ExtensionXsd = $@"{SourceHomeBuilderPath}\extensions\extension.xsd";

build/ArtifactBuilder/FrameworkAgentComponents.cs

@@ -67,6 +67,7 @@ protected override void CreateAgentComponents()
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.MassTransit.dll",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.MassTransitLegacy.dll",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Kafka.dll",
+                $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Bedrock.dll",
             };
 
             var wrapperXmls = new[]
@@ -105,6 +106,7 @@ protected override void CreateAgentComponents()
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.MassTransit.Instrumentation.xml",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.MassTransitLegacy.Instrumentation.xml",
                 $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Kafka.Instrumentation.xml",
+                $@"{SourceHomeBuilderPath}\extensions\NewRelic.Providers.Wrapper.Bedrock.Instrumentation.xml",
             };
 
             ExtensionXsd = $@"{SourceHomeBuilderPath}\extensions\extension.xsd";

src/Agent/MsiInstaller/Installer/Product.wxs

@@ -474,6 +474,9 @@ SPDX-License-Identifier: Apache-2.0
       <Component Id="KafkaWrapperComponent" Guid="{BE97F5F5-A392-48A0-8888-B2973EA09490}">
         <File Id="KafkaWrapperFile" Name="NewRelic.Providers.Wrapper.Kafka.dll" KeyPath="yes" Source="$(var.HomeFolderPath)\extensions\NewRelic.Providers.Wrapper.Kafka.dll"/>
       </Component>
+      <Component Id="BedrockWrapperComponent" Guid="{C913E912-97D1-4042-8E11-A35E04C6A6E5}">
+        <File Id="BedrockWrapperFile" Name="NewRelic.Providers.Wrapper.Bedrock.dll" KeyPath="yes" Source="$(var.HomeFolderPath)\extensions\NewRelic.Providers.Wrapper.Bedrock.dll"/>
+      </Component>
 
       <!-- Reference libraries -->
       <Component Id="NewRelicCoreReferenceComponent" Guid="{C196ED1E-0FBA-4D36-9C34-E969B0C9E8C9}">
@@ -545,6 +548,9 @@ SPDX-License-Identifier: Apache-2.0
       <Component Id="CoreAspNetCore6PlusWrapperComponent" Guid="{3B7C5E60-BA9D-4B1C-8B16-B16025B075C0}">
         <File Id="CoreAspNetCore6PlusWrapperFile" Name="NewRelic.Providers.Wrapper.AspNetCore6Plus.dll" KeyPath="yes" Source="$(var.HomeFolderPath)_coreclr\extensions\NewRelic.Providers.Wrapper.AspNetCore6Plus.dll"/>
       </Component>
+      <Component Id="CoreBedrockWrapperComponent" Guid="{7DA522F5-38D4-4E7E-AA2B-857D01812152}">
+        <File Id="CoreBedrockWrapperFile" Name="NewRelic.Providers.Wrapper.Bedrock.dll" KeyPath="yes" Source="$(var.HomeFolderPath)_coreclr\extensions\NewRelic.Providers.Wrapper.Bedrock.dll"/>
+      </Component>
 
       <!-- Reference libraries -->
       <Component Id="CoreNewRelicCoreReferenceComponent" Guid="{DD2BE979-7D4B-47EA-9FBE-F6B381D70E0B}">
@@ -659,6 +665,9 @@ SPDX-License-Identifier: Apache-2.0
       <Component Id="KafkaInstrumentationComponent" Guid="{9FC86A0E-CACD-4DA8-84F8-20997C904913}">
         <File Id="KafkaInstrumentationFile" Name="NewRelic.Providers.Wrapper.Kafka.Instrumentation.xml" KeyPath="yes" Source="$(var.HomeFolderPath)\extensions\NewRelic.Providers.Wrapper.Kafka.Instrumentation.xml"/>
       </Component>
+      <Component Id="BedrockInstrumentationComponent" Guid="{D6F0A2D6-D8D5-4C61-8056-1C9AA4D07132}">
+        <File Id="BedrockInstrumentationFile" Name="NewRelic.Providers.Wrapper.Bedrock.Instrumentation.xml" KeyPath="yes" Source="$(var.HomeFolderPath)\extensions\NewRelic.Providers.Wrapper.Bedrock.Instrumentation.xml"/>
+      </Component>
     </ComponentGroup>
 
     <ComponentGroup Id="CoreNewRelic.Agent.Extensions.Instrumentation" Directory="CoreExtensionsFolder">
@@ -719,6 +728,9 @@ SPDX-License-Identifier: Apache-2.0
       <Component Id="CoreAspNetCore6PlusInstrumentationComponent" Guid="{1CC1E672-5AA5-4F6F-A736-748EA556653A}">
         <File Id="CoreAspNetCore6PlusInstrumentationFile" Name="NewRelic.Providers.Wrapper.AspNetCore6Plus.Instrumentation.xml" KeyPath="yes" Source="$(var.HomeFolderPath)_coreclr\extensions\NewRelic.Providers.Wrapper.AspNetCore6Plus.Instrumentation.xml"/>
       </Component>
+      <Component Id="CoreBedrockInstrumentationComponent" Guid="{CF3A17F1-3C5C-4B72-8E3B-10E53EB94913}">
+        <File Id="CoreBedrockInstrumentationFile" Name="NewRelic.Providers.Wrapper.Bedrock.Instrumentation.xml" KeyPath="yes" Source="$(var.HomeFolderPath)_coreclr\extensions\NewRelic.Providers.Wrapper.Bedrock.Instrumentation.xml"/>
+      </Component>
     </ComponentGroup>
 
     <!-- Extensions XSD-->

src/Agent/NewRelic.Api.Agent/NewRelic.cs

@@ -849,6 +849,45 @@ public static void SetErrorGroupCallback(Func<IReadOnlyDictionary<string, object
             }
         }
 
+        /// <summary> Sets the method that will be invoked to define the token count of completion.
+        ///
+        /// The callback takes the model name and input value, and returns an integer of the token count.
+        /// A value returned from the callback that is less than or equal to 0 will be ignored.
+        /// </summary>
+        /// <param name="callback">The callback to invoke to generate the token count based on the model and input..</param>
+        [MethodImpl(MethodImplOptions.NoOptimization | MethodImplOptions.NoInlining)]
+        public static void SetLlmTokenCountingCallback(Func<string, string, int> callback)
+        {
+            try
+            {
+                System.Diagnostics.Trace.WriteLine("NewRelic.SetLlmTokenCountingCallback()");
+            }
+            catch
+            {
+                // Swallow any exception thrown from here
+            }
+        }
+
+        /// <summary>
+        /// Creates an event with the customer feedback on the LLM interaction.
+        /// </summary>
+        /// <param name="traceId">Required. ID of the trace where the chat completion(s) related to the feedback occurred</param>
+        /// <param name="rating">Required. Rating provided by an end user. Must be string or int</param>
+        /// <param name="category">Optional. Category of the feedback as provided by the end user</param>
+        /// <param name="message">Optional. Freeform text feedback from an end user</param>
+        /// <param name="metadata">Optional. Set of key-value pairs to store any other desired data to submit with the feedback event</param>
+        public static void RecordLlmFeedbackEvent(string traceId, object rating, string category = "", string message = "", IDictionary<string, object>? metadata = null)
+        {
+            try
+            {
+                System.Diagnostics.Trace.WriteLine("NewRelic.RecordLlmFeedbackEvent()");
+            }
+            catch
+            {
+                // Swallow any exception thrown from here
+            }
+        }
+
         #endregion
     }
 

src/Agent/NewRelic/Agent/Core/Agent.cs

@@ -13,6 +13,7 @@
 using NewRelic.Agent.Core.Metrics;
 using NewRelic.Agent.Core.Segments;
 using NewRelic.Agent.Core.Transactions;
+using NewRelic.Agent.Core.Transformers;
 using NewRelic.Agent.Core.Transformers.TransactionTransformer;
 using NewRelic.Agent.Core.Utilities;
 using NewRelic.Agent.Core.WireModels;
@@ -65,14 +66,17 @@ public class Agent : IAgent // any changes to api, update the interface in exten
         private Extensions.Logging.ILogger _logger;
         private readonly ISimpleSchedulingService _simpleSchedulingService;
 
+        private readonly ICustomEventTransformer _customEventTransformer;
+
         public Agent(ITransactionService transactionService, ITransactionTransformer transactionTransformer,
             IThreadPoolStatic threadPoolStatic, ITransactionMetricNameMaker transactionMetricNameMaker, IPathHashMaker pathHashMaker,
             ICatHeaderHandler catHeaderHandler, IDistributedTracePayloadHandler distributedTracePayloadHandler,
             ISyntheticsHeaderHandler syntheticsHeaderHandler, ITransactionFinalizer transactionFinalizer,
             IBrowserMonitoringPrereqChecker browserMonitoringPrereqChecker, IBrowserMonitoringScriptMaker browserMonitoringScriptMaker,
             IConfigurationService configurationService, IAgentHealthReporter agentHealthReporter, IAgentTimerService agentTimerService,
             IMetricNameService metricNameService, Api.ITraceMetadataFactory traceMetadataFactory, ICATSupportabilityMetricCounters catMetricCounters,
-            ILogEventAggregator logEventAggregator, ILogContextDataFilter logContextDataFilter, ISimpleSchedulingService simpleSchedulingService)
+            ILogEventAggregator logEventAggregator, ILogContextDataFilter logContextDataFilter, ISimpleSchedulingService simpleSchedulingService,
+            ICustomEventTransformer customEventTransformer)
         {
             _transactionService = transactionService;
             _transactionTransformer = transactionTransformer;
@@ -95,6 +99,8 @@ public Agent(ITransactionService transactionService, ITransactionTransformer tra
             _logContextDataFilter = logContextDataFilter;
             _simpleSchedulingService = simpleSchedulingService;
 
+            _customEventTransformer = customEventTransformer;
+
             Instance = this;
         }
 
@@ -411,6 +417,61 @@ public Dictionary<string, string> GetLinkingMetadata()
 
         #region ExperimentalApi
 
+        public void RecordLlmEvent(string eventType, IDictionary<string, object> attributes)
+        {
+            if (!_configurationService.Configuration.AiMonitoringEnabled)
+            {
+                return;
+            }
+
+            // Record metric is streaming has been disabled
+            if (!_configurationService.Configuration.AiMonitoringStreamingEnabled)
+            {
+                RecordSupportabilityMetric("Supportability/DotNet/ML/Streaming/Disabled");
+            }
+
+            var transaction = _transactionService.GetCurrentInternalTransaction();
+            transaction.SetLlmTransaction(true);
+
+            // Any custom attributes that are prefixed with "llm." must be added to the event
+            var transactionAttributes = transaction.TransactionMetadata.UserAndRequestAttributes.GetAllAttributeValuesDic();
+            foreach (var attribute in transactionAttributes)
+            {
+                if (attribute.Key.StartsWith("llm.", StringComparison.InvariantCultureIgnoreCase))
+                {
+                    attributes.Add(attribute.Key, attribute.Value);
+                }
+            }
+
+            // Always use callback for token count if one is provided
+            if ((eventType == "LlmChatCompletionMessage" || eventType == "LlmEmbedding")
+                && _configurationService.Configuration.LlmTokenCountingCallback != null)
+            {
+                // message and embedding events have different attribute names for the content of the message
+                var content = eventType == "LlmChatCompletionMessage" ? (string)attributes["content"] : (string)attributes["input"];
+
+                // messages only have response models, embeddings have both
+                var model = attributes.TryGetValue("request.model", out var requestModel) ? (string)requestModel : (string)attributes["response.model"];
+
+                // Use a nullable so that the CustomEvent code will automatically remove the attribute if the callback returns null
+                var tokenCount = _configurationService.Configuration.LlmTokenCountingCallback?.Invoke(model, content);
+                if (tokenCount.HasValue && tokenCount.Value > 0)
+                {
+                    attributes["token_count"] = tokenCount;
+                }
+            }
+
+            // If HSM is enabled or if record content is disabled, we need to remove the content and input attributes
+            // We will still want the token counts so removal occurs after the attempt to get the token count
+            if (!_configurationService.Configuration.AiMonitoringRecordContentEnabled || _configurationService.Configuration.HighSecurityModeEnabled)
+            {
+                attributes.Remove("content"); // ChatMessages
+                attributes.Remove("input"); // Embeddings
+            }
+
+            _customEventTransformer.Transform(eventType, attributes, transaction.Priority);
+        }
+
         public ISimpleSchedulingService SimpleSchedulingService
         {
             get { return _simpleSchedulingService; }