Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

summary: AWS Bedrock instrumentation #2314

Merged
merged 43 commits into from
Apr 2, 2024
Merged
Show file tree
Hide file tree
Changes from 39 commits
Commits
Show all changes
43 commits
Select commit Hold shift + click to select a range
e28e4e1
Most AI models and working events.
jaffinito Feb 8, 2024
4d87904
Adds Cohere Command and fixes version metric
jaffinito Feb 8, 2024
e8658d6
Add support for multiple response and count those tokens.
jaffinito Feb 9, 2024
570306e
Added commented section for Titan embed - the only one there is
jaffinito Feb 9, 2024
4c50a68
Correct operation metric.
jaffinito Feb 9, 2024
677ff66
Only create events in transcations
jaffinito Feb 9, 2024
7f45817
PoC for milestone with real types
jaffinito Feb 13, 2024
27860e5
Fixes Txn.Guid intergace problem.
jaffinito Feb 13, 2024
091d830
Massive rework of PoC to align with merged specs.
jaffinito Feb 28, 2024
ccb6822
Adds record content config and HSM support.
jaffinito Feb 28, 2024
4e88298
Remove public TransactionId
jaffinito Feb 28, 2024
3c53674
Unwind new TransactionId API
jaffinito Feb 28, 2024
4f8fbc9
Unwind new TransactionId API 2
jaffinito Feb 28, 2024
a38117f
Feedback API
jaffinito Feb 28, 2024
7b065b9
Add missing supportability metric
jaffinito Feb 29, 2024
e967be4
Merge poc branch to feature branch, resolve conflicts
tippmar-nr Mar 4, 2024
12e96c0
Add Bedrock wrapper as a dependency of the Home project
tippmar-nr Mar 4, 2024
3f18c99
Refactoring, updated unit tests to pass with new ai supportability me…
tippmar-nr Mar 4, 2024
35b19a3
Update ArtifactBuilder / MSI to add Bedrock to the extensions folder
tippmar-nr Mar 5, 2024
e3c1423
Remove reference to AWSSDK.BedrockRuntime, convert variables to dynam…
tippmar-nr Mar 5, 2024
3199238
Implemented error handling for Bedrock (#2312)
tippmar-nr Mar 7, 2024
5c78cbe
Merge remote-tracking branch 'origin/main' into feature/aws-bedrock
tippmar-nr Mar 8, 2024
a019002
Bedrock unit tests (#2321)
tippmar-nr Mar 8, 2024
41f842f
Merge latest from main, resolve conflicts
tippmar-nr Mar 11, 2024
f175e16
Fixing a merge mistake
tippmar-nr Mar 11, 2024
58c5660
Minor tweak
tippmar-nr Mar 11, 2024
b129778
Bedrock - Content Recording (#2333)
tippmar-nr Mar 12, 2024
846ca77
Merge remote-tracking branch 'origin/main' into feature/aws-bedrock
tippmar-nr Mar 13, 2024
edb5c82
Bedrock: Revert some work, add some unit tests (#2339)
tippmar-nr Mar 13, 2024
67e174f
Bedrock integration tests (#2340)
chynesNR Mar 15, 2024
ac7aacc
Restored the copyright header
tippmar-nr Mar 19, 2024
ade65a0
Merge branch 'feature/aws-bedrock' of https://github.com/newrelic/new…
tippmar-nr Mar 19, 2024
828a585
Cleanup unused usings and local var (#2352)
nr-ahemsath Mar 19, 2024
3834c09
Merge remote-tracking branch 'origin/main' into feature/aws-bedrock
tippmar-nr Mar 20, 2024
b0e3b2c
Remove reference to `System.Text.Json`, refactor to use `Newtonsoft.J…
tippmar-nr Mar 20, 2024
a13c268
Log message when response model is null or non-200 status code
tippmar-nr Mar 20, 2024
679235d
Fix request/response payload parsing issues after converting to Newto…
tippmar-nr Mar 21, 2024
6675f08
Refactoring and tests for RecordLlmFeedbackEvent (#2354)
tippmar-nr Mar 21, 2024
1aa285d
Disable AI monitoring when HSM is enabled (#2353)
tippmar-nr Mar 25, 2024
cb8d1ae
Merge latest from main, resolve conflicts
tippmar-nr Mar 26, 2024
73dea59
Move JSON deserialization and models to NewRelic.Core (#2378)
nr-ahemsath Mar 26, 2024
f030e84
Fix population of `role` attribute and remove `tokenCount` attribute …
tippmar-nr Mar 28, 2024
ee4cd31
Fix issue with errors not appearing in the UI (#2382)
chynesNR Apr 1, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/all_solutions.yml
Original file line number Diff line number Diff line change
Expand Up @@ -247,7 +247,8 @@ jobs:
DistributedTracing,
Errors,
HttpClientInstrumentation,
InfiniteTracing,
InfiniteTracing,
LLM,
Logging.ContextData,
Logging.HsmAndCsp,
Logging.LocalDecoration,
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/run_integration_tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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', \
Expand Down
8 changes: 8 additions & 0 deletions FullAgent.sln
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions build/ArtifactBuilder/CoreAgentComponents.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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[]
Expand All @@ -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";
Expand Down
2 changes: 2 additions & 0 deletions build/ArtifactBuilder/FrameworkAgentComponents.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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[]
Expand Down Expand Up @@ -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";
Expand Down
12 changes: 12 additions & 0 deletions src/Agent/MsiInstaller/Installer/Product.wxs
Original file line number Diff line number Diff line change
Expand Up @@ -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}">
Expand Down Expand Up @@ -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}">
Expand Down Expand Up @@ -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">
Expand Down Expand Up @@ -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-->
Expand Down
39 changes: 39 additions & 0 deletions src/Agent/NewRelic.Api.Agent/NewRelic.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
63 changes: 62 additions & 1 deletion src/Agent/NewRelic/Agent/Core/Agent.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -95,6 +99,8 @@ public Agent(ITransactionService transactionService, ITransactionTransformer tra
_logContextDataFilter = logContextDataFilter;
_simpleSchedulingService = simpleSchedulingService;

_customEventTransformer = customEventTransformer;

Instance = this;
}

Expand Down Expand Up @@ -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 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)
{
attributes.Remove("content"); // ChatMessages
attributes.Remove("input"); // Embeddings
}

_customEventTransformer.Transform(eventType, attributes, transaction.Priority);
}

public ISimpleSchedulingService SimpleSchedulingService
{
get { return _simpleSchedulingService; }
Expand Down
34 changes: 34 additions & 0 deletions src/Agent/NewRelic/Agent/Core/Api/AgentApi.cs
Original file line number Diff line number Diff line change
Expand Up @@ -642,6 +642,40 @@ void work()
}
TryInvoke(work, apiName, ApiMethod.SetErrorGroupCallback);
}

/// <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>
public static void SetLlmTokenCountingCallback(Func<string, string, int> callback)
{
const string apiName = nameof(SetLlmTokenCountingCallback);
void work()
{
InternalApi.SetLlmTokenCountingCallback(callback);
}
TryInvoke(work, apiName, ApiMethod.SetLlmTokenCountingCallback);
}

/// <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)
{
const string apiName = nameof(RecordLlmFeedbackEvent);
void work()
{
InternalApi.RecordLlmFeedbackEvent(traceId, rating, category, message, metadata);
}
TryInvoke(work, apiName, ApiMethod.RecordLlmFeedbackEvent);
}
}
}

Expand Down
Loading
Loading