Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
113 changes: 98 additions & 15 deletions DevProxy.Plugins/Inspection/OpenAITelemetryPlugin.cs
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,12 @@
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using OpenTelemetry;
using OpenTelemetry.Exporter;
using OpenTelemetry.Metrics;
using OpenTelemetry.Resources;
using OpenTelemetry.Trace;
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Diagnostics.Metrics;
using System.Text.Json;
Expand Down Expand Up @@ -67,7 +67,6 @@ public sealed class OpenAITelemetryPlugin(
private LanguageModelPricesLoader? _loader;
private MeterProvider? _meterProvider;
private TracerProvider? _tracerProvider;
private readonly ConcurrentDictionary<string, List<OpenAITelemetryPluginReportModelUsageInformation>> _modelUsage = [];

public override string Name => nameof(OpenAITelemetryPlugin);

Expand Down Expand Up @@ -196,17 +195,18 @@ public override Task AfterRecordingStopAsync(RecordingArgs e, CancellationToken
{
Logger.LogTrace("{Method} called", nameof(AfterRecordingStopAsync));

ArgumentNullException.ThrowIfNull(e);

var report = new OpenAITelemetryPluginReport
{
Application = Configuration.Application,
Environment = Configuration.Environment,
Currency = Configuration.Currency,
IncludeCosts = Configuration.IncludeCosts,
ModelUsage = _modelUsage.ToDictionary()
ModelUsage = GetOpenAIModelUsage(e.RequestLogs)
};

StoreReport(report, e);
_modelUsage.Clear();

Logger.LogTrace("Left {Name}", nameof(AfterRecordingStopAsync));
return Task.CompletedTask;
Expand Down Expand Up @@ -849,16 +849,6 @@ private void RecordUsageMetrics(Activity activity, OpenAIRequest request, OpenAI
.SetTag(SemanticConvention.GEN_AI_USAGE_OUTPUT_TOKENS, usage.CompletionTokens)
.SetTag(SemanticConvention.GEN_AI_USAGE_TOTAL_TOKENS, usage.TotalTokens);

var reportModelUsageInformation = new OpenAITelemetryPluginReportModelUsageInformation
{
Model = response.Model,
PromptTokens = usage.PromptTokens,
CompletionTokens = usage.CompletionTokens,
CachedTokens = usage.PromptTokensDetails?.CachedTokens ?? 0L
};
var usagePerModel = _modelUsage.GetOrAdd(response.Model, model => []);
usagePerModel.Add(reportModelUsageInformation);

if (!Configuration.IncludeCosts || Configuration.Prices is null)
{
Logger.LogDebug("Cost tracking is disabled or prices data is not available");
Expand Down Expand Up @@ -895,7 +885,6 @@ private void RecordUsageMetrics(Activity activity, OpenAIRequest request, OpenAI
new(SemanticConvention.GEN_AI_REQUEST_MODEL, request.Model),
new(SemanticConvention.GEN_AI_RESPONSE_MODEL, response.Model)
]);
reportModelUsageInformation.Cost = totalCost;
}
else
{
Expand All @@ -905,6 +894,100 @@ private void RecordUsageMetrics(Activity activity, OpenAIRequest request, OpenAI
Logger.LogTrace("RecordUsageMetrics() finished");
}

private Dictionary<string, List<OpenAITelemetryPluginReportModelUsageInformation>> GetOpenAIModelUsage(IEnumerable<RequestLog> requestLogs)
{
var modelUsage = new Dictionary<string, List<OpenAITelemetryPluginReportModelUsageInformation>>();
var openAIRequestLogs = requestLogs.Where(r =>
r is not null &&
r.Context is not null &&
r.Context.Session is not null &&
r.MessageType == MessageType.InterceptedResponse &&
string.Equals("POST", r.Context.Session.HttpClient.Request.Method, StringComparison.OrdinalIgnoreCase) &&
r.Context.Session.HttpClient.Response.StatusCode >= 200 &&
r.Context.Session.HttpClient.Response.StatusCode < 300 &&
r.Context.Session.HttpClient.Response.HasBody &&
!string.IsNullOrEmpty(r.Context.Session.HttpClient.Response.BodyString) &&
ProxyUtils.MatchesUrlToWatch(UrlsToWatch, r.Context.Session.HttpClient.Request.RequestUri.AbsoluteUri) &&
OpenAIRequest.TryGetOpenAIRequest(r.Context.Session.HttpClient.Request.BodyString, NullLogger.Instance, out var openAiRequest) &&
openAiRequest is not null
);

foreach (var requestLog in openAIRequestLogs)
{
try
{
var response = JsonSerializer.Deserialize<OpenAIResponse>(requestLog.Context!.Session.HttpClient.Response.BodyString, ProxyUtils.JsonSerializerOptions);
if (response is null)
{
continue;
}

var reportModelUsageInfo = GetReportModelUsageInfo(response);
if (modelUsage.TryGetValue(response.Model, out var usagePerModel))
{
usagePerModel.AddRange(reportModelUsageInfo);
}
else
{
modelUsage.Add(response.Model, reportModelUsageInfo);
}
}
catch (JsonException ex)
{
Logger.LogError(ex, "Failed to deserialize OpenAI response");
}
}

return modelUsage;
}

private List<OpenAITelemetryPluginReportModelUsageInformation> GetReportModelUsageInfo(OpenAIResponse response)
{
Logger.LogTrace("GetReportModelUsageInfo() called");
var usagePerModel = new List<OpenAITelemetryPluginReportModelUsageInformation>();
var usage = response.Usage;
if (usage is null)
{
return usagePerModel;
}

var reportModelUsageInformation = new OpenAITelemetryPluginReportModelUsageInformation
{
Model = response.Model,
PromptTokens = usage.PromptTokens,
CompletionTokens = usage.CompletionTokens,
CachedTokens = usage.PromptTokensDetails?.CachedTokens ?? 0L
};
usagePerModel.Add(reportModelUsageInformation);

if (!Configuration.IncludeCosts || Configuration.Prices is null)
{
Logger.LogDebug("Cost tracking is disabled or prices data is not available");
return usagePerModel;
}

if (string.IsNullOrEmpty(response.Model))
{
Logger.LogDebug("Response model is empty or null");
return usagePerModel;
}

var (inputCost, outputCost) = Configuration.Prices.CalculateCost(response.Model, usage.PromptTokens, usage.CompletionTokens);

if (inputCost > 0)
{
var totalCost = inputCost + outputCost;
reportModelUsageInformation.Cost = totalCost;
}
else
{
Logger.LogDebug("Input cost is zero, skipping cost metrics recording");
}

Logger.LogTrace("GetReportModelUsageInfo() finished");
return usagePerModel;
}

private static string GetOperationName(OpenAIRequest request)
{
if (request == null)
Expand Down
2 changes: 1 addition & 1 deletion DevProxy/Proxy/ProxyStateController.cs
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ public async Task StopRecordingAsync(CancellationToken cancellationToken)

public async Task MockRequestAsync(CancellationToken cancellationToken)
{
var eventArgs = new EventArgs();
var eventArgs = EventArgs.Empty;

foreach (var plugin in _plugins.Where(p => p.Enabled))
{
Expand Down