From 28e1a41f382afa9d7ead061213329b6c0b36082e Mon Sep 17 00:00:00 2001 From: Brendan Gooden <30429990+brendangooden@users.noreply.github.com> Date: Tue, 14 Apr 2026 11:17:05 +1200 Subject: [PATCH 1/2] feat: Add OTEL GenAI attributes to tool spans and raw SDK packages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add gen_ai.* attributes where Braintrust fills gaps no framework covers: - Function/tool spans: gen_ai.operation.name ("execute_tool"), gen_ai.tool.name, gen_ai.tool.call.arguments, gen_ai.tool.call.result (M.E.AI's UseOpenTelemetry() does not emit dedicated execute_tool spans) - OpenAI raw SDK: gen_ai.operation.name, gen_ai.provider.name, gen_ai.input/output.messages, gen_ai.usage.input/output_tokens (no framework-level OTEL instrumentation for raw ChatClient) - Anthropic raw SDK: same gen_ai.* attributes (no framework-level OTEL instrumentation for raw IMessageService) IChatClient chat/agent spans intentionally omit gen_ai.* — users should stack UseOpenTelemetry() from M.E.AI for standard OTEL GenAI attributes. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../BraintrustFunctionMiddleware.cs | 12 ++++++---- .../InstrumentedMessageService.cs | 24 +++++++++++++------ .../InstrumentedChatClient.cs | 22 ++++++++++++++--- .../FunctionMiddlewareTests.cs | 11 +++++++++ 4 files changed, 55 insertions(+), 14 deletions(-) diff --git a/src/Braintrust.Sdk.AgentFramework/BraintrustFunctionMiddleware.cs b/src/Braintrust.Sdk.AgentFramework/BraintrustFunctionMiddleware.cs index 105ee9c..1bf30ad 100644 --- a/src/Braintrust.Sdk.AgentFramework/BraintrustFunctionMiddleware.cs +++ b/src/Braintrust.Sdk.AgentFramework/BraintrustFunctionMiddleware.cs @@ -35,7 +35,9 @@ internal static class BraintrustFunctionMiddleware if (activity != null) { SpanTagHelper.SetSpanType(activity, "function_call"); + activity.SetTag("gen_ai.operation.name", "execute_tool"); activity.SetTag("function.name", functionName); + activity.SetTag("gen_ai.tool.name", functionName); activity.SetTag("function.iteration", context.Iteration); activity.SetTag("function.call_index", context.FunctionCallIndex); activity.SetTag("function.total_count", context.FunctionCount); @@ -44,8 +46,9 @@ internal static class BraintrustFunctionMiddleware { try { - activity.SetTag("braintrust.input_json", - SpanTagHelper.ToJson(context.Arguments)); + var argsJson = SpanTagHelper.ToJson(context.Arguments); + activity.SetTag("braintrust.input_json", argsJson); + activity.SetTag("gen_ai.tool.call.arguments", argsJson); } catch { @@ -75,8 +78,9 @@ internal static class BraintrustFunctionMiddleware { try { - activity.SetTag("braintrust.output_json", - SpanTagHelper.ToJson(new { result })); + var resultJson = SpanTagHelper.ToJson(new { result }); + activity.SetTag("braintrust.output_json", resultJson); + activity.SetTag("gen_ai.tool.call.result", resultJson); } catch { diff --git a/src/Braintrust.Sdk.Anthropic/InstrumentedMessageService.cs b/src/Braintrust.Sdk.Anthropic/InstrumentedMessageService.cs index 76498e7..b8577be 100644 --- a/src/Braintrust.Sdk.Anthropic/InstrumentedMessageService.cs +++ b/src/Braintrust.Sdk.Anthropic/InstrumentedMessageService.cs @@ -166,12 +166,12 @@ public async IAsyncEnumerable CreateStreaming( activity.SetTag("braintrust.metrics.time_to_first_token", timeToFirstToken.Value); } - activity.SetTag( - "braintrust.output_json", - ToJson(new object[] - { - new { role = role ?? "assistant", content = output.ToString() } - })); + var outputJson = ToJson(new object[] + { + new { role = role ?? "assistant", content = output.ToString() } + }); + activity.SetTag("braintrust.output_json", outputJson); + activity.SetTag("gen_ai.output.messages", outputJson); } } @@ -203,6 +203,8 @@ private static void TagActivity( double? timeToFirstToken = null) { activity.SetTag("provider", "anthropic"); + activity.SetTag("gen_ai.operation.name", "chat"); + activity.SetTag("gen_ai.provider.name", "anthropic"); activity.SetTag("gen_ai.request.model", request.Model.Raw()); activity.SetTag("gen_ai.response.model", response.Model.Raw()); @@ -221,14 +223,19 @@ private static void TagActivity( sys.TryPickString(out var sysContent); inputMessages.Add(new { role = "system", content = sysContent }); } - activity.SetTag("braintrust.input_json", ToJson(inputMessages)); + var inputJson = ToJson(inputMessages); + activity.SetTag("braintrust.input_json", inputJson); + activity.SetTag("gen_ai.input.messages", inputJson); var contentJson = response.ToString(); activity.SetTag("braintrust.output_json", contentJson); + activity.SetTag("gen_ai.output.messages", contentJson); // Extract token usage metrics activity.SetTag("braintrust.metrics.prompt_tokens", response.Usage.InputTokens); + activity.SetTag("gen_ai.usage.input_tokens", response.Usage.InputTokens); activity.SetTag("braintrust.metrics.completion_tokens", response.Usage.OutputTokens); + activity.SetTag("gen_ai.usage.output_tokens", response.Usage.OutputTokens); activity.SetTag("braintrust.metrics.tokens", response.Usage.InputTokens + response.Usage.OutputTokens); if (timeToFirstToken is > 0) @@ -259,12 +266,15 @@ private static void TagStreamActivity(Activity activity, MessageCreateParams req { activity.SetTag("stream", true); activity.SetTag("provider", "anthropic"); + activity.SetTag("gen_ai.operation.name", "chat"); + activity.SetTag("gen_ai.provider.name", "anthropic"); activity.SetTag("gen_ai.request.model", request.Model.Raw()); try { var messagesJson = ToJson(request.Messages); activity.SetTag("braintrust.input_json", messagesJson); + activity.SetTag("gen_ai.input.messages", messagesJson); } catch { diff --git a/src/Braintrust.Sdk.OpenAI/InstrumentedChatClient.cs b/src/Braintrust.Sdk.OpenAI/InstrumentedChatClient.cs index f45a4ea..066bbd5 100644 --- a/src/Braintrust.Sdk.OpenAI/InstrumentedChatClient.cs +++ b/src/Braintrust.Sdk.OpenAI/InstrumentedChatClient.cs @@ -125,17 +125,25 @@ public override async Task> CompleteChatAsync(IEnum } } - // TODO: Override other methods as needed (CompleteChatStreaming, etc.) + // Note: Streaming (CompleteChatStreaming/CompleteChatStreamingAsync) is best instrumented + // via the IChatClient adapter path: chatClient.AsIChatClient() + Braintrust.Sdk.Extensions.AI. + // The raw OpenAI ChatClient streaming API returns AsyncCollectionResult which cannot be + // transparently wrapped without breaking the SSE response stream. + private void TagActivity(Activity activity, double? timeToFirstToken = null) { activity.SetTag("provider", "openai"); + activity.SetTag("gen_ai.operation.name", "chat"); + activity.SetTag("gen_ai.provider.name", "openai"); { var requestRaw = activity.GetBaggageItem("braintrust.http.request"); if (requestRaw != null) { var requestJson = JsonNode.Parse(requestRaw); activity.SetTag("gen_ai.request.model", requestJson?["model"]?.ToString()); - activity.SetTag("braintrust.input_json", requestJson?["messages"]?.ToString()); + var messagesJson = requestJson?["messages"]?.ToString(); + activity.SetTag("braintrust.input_json", messagesJson); + activity.SetTag("gen_ai.input.messages", messagesJson); } } { @@ -144,7 +152,9 @@ private void TagActivity(Activity activity, double? timeToFirstToken = null) { var responseJson = JsonNode.Parse(responseRaw); activity.SetTag("gen_ai.response.model", responseJson?["model"]?.ToString()); - activity.SetTag("braintrust.output_json", responseJson?["choices"]?.ToString()); + var choicesJson = responseJson?["choices"]?.ToString(); + activity.SetTag("braintrust.output_json", choicesJson); + activity.SetTag("gen_ai.output.messages", choicesJson); // Extract token usage metrics var usage = responseJson?["usage"]; @@ -155,9 +165,15 @@ private void TagActivity(Activity activity, double? timeToFirstToken = null) var totalTokens = usage["total_tokens"]?.GetValue(); if (promptTokens.HasValue) + { activity.SetTag("braintrust.metrics.prompt_tokens", promptTokens.Value); + activity.SetTag("gen_ai.usage.input_tokens", promptTokens.Value); + } if (completionTokens.HasValue) + { activity.SetTag("braintrust.metrics.completion_tokens", completionTokens.Value); + activity.SetTag("gen_ai.usage.output_tokens", completionTokens.Value); + } if (totalTokens.HasValue) activity.SetTag("braintrust.metrics.tokens", totalTokens.Value); } diff --git a/tests/Braintrust.Sdk.AgentFramework.Tests/FunctionMiddlewareTests.cs b/tests/Braintrust.Sdk.AgentFramework.Tests/FunctionMiddlewareTests.cs index dfcd6bc..29dc9cb 100644 --- a/tests/Braintrust.Sdk.AgentFramework.Tests/FunctionMiddlewareTests.cs +++ b/tests/Braintrust.Sdk.AgentFramework.Tests/FunctionMiddlewareTests.cs @@ -47,6 +47,10 @@ public async Task UseBraintrustFunctionTracing_CreatesSpanForFunctionCall() var spanType = funcActivity.GetTagItem("braintrust.span_attributes")?.ToString(); Assert.Contains("\"type\":\"function_call\"", spanType); Assert.Equal("GetWeather", funcActivity.GetTagItem("function.name")); + + // OTEL GenAI tool span attributes (Braintrust value-add — M.E.AI doesn't emit these) + Assert.Equal("execute_tool", funcActivity.GetTagItem("gen_ai.operation.name")); + Assert.Equal("GetWeather", funcActivity.GetTagItem("gen_ai.tool.name")); } [Fact] @@ -146,6 +150,13 @@ public async Task UseBraintrustFunctionTracing_CapturesArgumentsAndResult() var outputJson = funcActivity.GetTagItem("braintrust.output_json")?.ToString(); Assert.NotNull(outputJson); Assert.Contains("Sunny", outputJson); + + // OTEL GenAI tool call attributes (Braintrust value-add) + var genAiArgs = funcActivity.GetTagItem("gen_ai.tool.call.arguments")?.ToString(); + Assert.NotNull(genAiArgs); + var genAiResult = funcActivity.GetTagItem("gen_ai.tool.call.result")?.ToString(); + Assert.NotNull(genAiResult); + Assert.Contains("Sunny", genAiResult); } } From ce5538f2e5a4d39a7f3a47ef3aac74b7b24c811f Mon Sep 17 00:00:00 2001 From: Brendan Gooden <30429990+brendangooden@users.noreply.github.com> Date: Tue, 14 Apr 2026 11:17:20 +1200 Subject: [PATCH 2/2] feat: Add Braintrust.Sdk.Extensions.AI package for IChatClient instrumentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New package that instruments any IChatClient implementation via Microsoft.Extensions.AI middleware pattern. Works with OpenAI, Azure OpenAI, Ollama, and any other IChatClient provider. Emits only braintrust.* attributes on chat/LLM spans (users should stack UseOpenTelemetry() for standard gen_ai.* attributes). Emits gen_ai.tool.* on execute_tool spans — the gap M.E.AI doesn't cover. Includes: - BraintrustChatClient: delegating IChatClient with LLM tracing - BraintrustFunctionMiddleware: tool/function call tracing - Extension methods: UseBraintrustTracing, UseBraintrustFunctionTracing, UseAllBraintrustTracing - Full test suite (6 tests) - Example project (ExtensionsAIInstrumentation) - Updated README with package selection guidance Also adds SkipVersionVerification property to Directory.Build.targets. Co-Authored-By: Claude Opus 4.6 (1M context) --- Directory.Build.targets | 2 +- README.md | 19 +- .../ExtensionsAIInstrumentation.csproj | 20 ++ .../ExtensionsAIInstrumentation/Program.cs | 72 +++++ .../Braintrust.Sdk.Extensions.AI.csproj | 27 ++ .../BraintrustChatClient.cs | 256 ++++++++++++++++ .../BraintrustExtensionsAI.cs | 103 +++++++ .../BraintrustFunctionMiddleware.cs | 119 ++++++++ src/Braintrust.Sdk.Extensions.AI/README.md | 16 + .../Braintrust.Sdk.Extensions.AI.Tests.csproj | 24 ++ .../BraintrustChatClientTests.cs | 275 ++++++++++++++++++ 11 files changed, 928 insertions(+), 5 deletions(-) create mode 100644 examples/ExtensionsAIInstrumentation/ExtensionsAIInstrumentation.csproj create mode 100644 examples/ExtensionsAIInstrumentation/Program.cs create mode 100644 src/Braintrust.Sdk.Extensions.AI/Braintrust.Sdk.Extensions.AI.csproj create mode 100644 src/Braintrust.Sdk.Extensions.AI/BraintrustChatClient.cs create mode 100644 src/Braintrust.Sdk.Extensions.AI/BraintrustExtensionsAI.cs create mode 100644 src/Braintrust.Sdk.Extensions.AI/BraintrustFunctionMiddleware.cs create mode 100644 src/Braintrust.Sdk.Extensions.AI/README.md create mode 100644 tests/Braintrust.Sdk.Extensions.AI.Tests/Braintrust.Sdk.Extensions.AI.Tests.csproj create mode 100644 tests/Braintrust.Sdk.Extensions.AI.Tests/BraintrustChatClientTests.cs diff --git a/Directory.Build.targets b/Directory.Build.targets index 8909c79..3d4d8cc 100644 --- a/Directory.Build.targets +++ b/Directory.Build.targets @@ -94,7 +94,7 @@ - + diff --git a/README.md b/README.md index a07944f..9800546 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ This library provides tools for **evaluating** and **tracing** AI applications i - **Evaluate** your AI models with custom test cases and scoring functions - **Trace** LLM calls and monitor AI application performance with OpenTelemetry -- **Integrate** seamlessly with OpenAI, Anthropic, Microsoft Agent Framework, and other LLM providers +- **Integrate** seamlessly with OpenAI, Anthropic, Microsoft Agent Framework, Microsoft.Extensions.AI (IChatClient), and other LLM providers This SDK is currently in BETA status and APIs may change. @@ -35,20 +35,31 @@ dotnet add package Braintrust.Sdk.OpenAI dotnet add package Braintrust.Sdk.Anthropic ``` +### Microsoft.Extensions.AI integration (IChatClient) + +```bash +dotnet add package Braintrust.Sdk.Extensions.AI +``` + +Works with any `IChatClient` provider: OpenAI, Azure OpenAI, Ollama, etc. + ### Microsoft Agent Framework integration ```bash dotnet add package Braintrust.Sdk.AgentFramework ``` +For agent orchestration with `ChatClientAgent`. Includes IChatClient + agent-level tracing. + ### Or add to your .csproj file ```xml - - - + + + + ``` diff --git a/examples/ExtensionsAIInstrumentation/ExtensionsAIInstrumentation.csproj b/examples/ExtensionsAIInstrumentation/ExtensionsAIInstrumentation.csproj new file mode 100644 index 0000000..7478adc --- /dev/null +++ b/examples/ExtensionsAIInstrumentation/ExtensionsAIInstrumentation.csproj @@ -0,0 +1,20 @@ + + + + Exe + net8.0 + enable + enable + + + + + + + + + + + + + diff --git a/examples/ExtensionsAIInstrumentation/Program.cs b/examples/ExtensionsAIInstrumentation/Program.cs new file mode 100644 index 0000000..9bd869f --- /dev/null +++ b/examples/ExtensionsAIInstrumentation/Program.cs @@ -0,0 +1,72 @@ +using Braintrust.Sdk; +using Braintrust.Sdk.Extensions.AI; +using Microsoft.Extensions.AI; + +namespace Braintrust.Sdk.Examples.ExtensionsAIInstrumentation; + +/// +/// Example demonstrating Braintrust instrumentation via Microsoft.Extensions.AI IChatClient. +/// Works with any provider: OpenAI, Azure OpenAI, Ollama, etc. +/// +/// Spans emitted include both braintrust.* (for Braintrust dashboard) and gen_ai.* +/// (OTEL GenAI semantic conventions) attributes. +/// +class Program +{ + static async Task Main(string[] args) + { + var openAIApiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY"); + if (string.IsNullOrEmpty(openAIApiKey)) + { + Console.WriteLine("ERROR: OPENAI_API_KEY environment variable not set. Bailing."); + return; + } + + var braintrust = Braintrust.Get(); + var activitySource = braintrust.GetActivitySource(); + + // Create an IChatClient from any provider — here using OpenAI via M.E.AI adapter + var openAIClient = new OpenAI.OpenAIClient(openAIApiKey); + IChatClient chatClient = openAIClient.GetChatClient("gpt-4o-mini").AsIChatClient(); + + // Add Braintrust tracing at both LLM and function levels + var tracedClient = new ChatClientBuilder(chatClient) + .UseAllBraintrustTracing(activitySource) + .Build(); + + // Define a tool + var getWeather = AIFunctionFactory.Create( + (string city) => $"The weather in {city} is sunny, 72°F.", + "GetWeather", + "Gets the current weather for a city."); + + using (var rootActivity = activitySource.StartActivity("extensions-ai-instrumentation-example")) + { + if (rootActivity != null) + { + Console.WriteLine("~~~ EXTENSIONS.AI INSTRUMENTATION EXAMPLE\n"); + + // Non-streaming call with tool use + var response = await tracedClient.GetResponseAsync( + [new ChatMessage(ChatRole.User, "What's the weather like in Seattle?")], + new ChatOptions { Tools = [getWeather] }); + + Console.WriteLine($"Response: {response.Text}"); + + // Streaming call + Console.Write("\nStreaming: "); + await foreach (var update in tracedClient.GetStreamingResponseAsync( + [new ChatMessage(ChatRole.User, "Tell me a joke.")])) + { + Console.Write(update.Text); + } + Console.WriteLine(); + + // Print Braintrust link + var url = await braintrust.GetProjectUriAsync() + + $"/logs?r={rootActivity.TraceId}&s={rootActivity.SpanId}"; + Console.WriteLine($"\n View your trace in Braintrust: {url}\n"); + } + } + } +} diff --git a/src/Braintrust.Sdk.Extensions.AI/Braintrust.Sdk.Extensions.AI.csproj b/src/Braintrust.Sdk.Extensions.AI/Braintrust.Sdk.Extensions.AI.csproj new file mode 100644 index 0000000..130f195 --- /dev/null +++ b/src/Braintrust.Sdk.Extensions.AI/Braintrust.Sdk.Extensions.AI.csproj @@ -0,0 +1,27 @@ + + + + net8.0 + enable + enable + + README.md + + + + + + + + + + + + + + + + + + + diff --git a/src/Braintrust.Sdk.Extensions.AI/BraintrustChatClient.cs b/src/Braintrust.Sdk.Extensions.AI/BraintrustChatClient.cs new file mode 100644 index 0000000..21b3bad --- /dev/null +++ b/src/Braintrust.Sdk.Extensions.AI/BraintrustChatClient.cs @@ -0,0 +1,256 @@ +using System.Diagnostics; +using System.Runtime.CompilerServices; +using System.Text; +using System.Text.Json; +using Microsoft.Extensions.AI; +using OpenTelemetry.Trace; + +namespace Braintrust.Sdk.Extensions.AI; + +/// +/// IChatClient middleware that wraps LLM calls with Braintrust tracing spans. +/// Captures prompts, completions, token usage, and timing metrics. +/// Emits braintrust.* attributes for Braintrust dashboard rendering. +/// For standard gen_ai.* OTEL attributes, use UseOpenTelemetry() from M.E.AI. +/// +internal sealed class BraintrustChatClient : DelegatingChatClient +{ + private readonly ActivitySource _activitySource; + private readonly bool _captureMessageContent; + + private static readonly JsonSerializerOptions JsonOptions = new() + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase, + DefaultIgnoreCondition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull + }; + + internal BraintrustChatClient( + IChatClient innerClient, + ActivitySource activitySource, + bool captureMessageContent) + : base(innerClient) + { + _activitySource = activitySource ?? throw new ArgumentNullException(nameof(activitySource)); + _captureMessageContent = captureMessageContent; + } + + public override async Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + using var activity = _activitySource.StartActivity("Chat Completion", ActivityKind.Client); + var startTime = DateTime.UtcNow; + + try + { + if (activity != null) + { + SetSpanType(activity, "llm"); + SetModel(activity, options?.ModelId); + + if (_captureMessageContent) + { + SetInputMessages(activity, messages); + } + } + + var response = await base.GetResponseAsync(messages, options, cancellationToken) + .ConfigureAwait(false); + + if (activity != null) + { + var timeToFirstToken = (DateTime.UtcNow - startTime).TotalSeconds; + + SetResponseModel(activity, response.ModelId); + SetTokenMetrics(activity, response.Usage); + SetTimeToFirstToken(activity, timeToFirstToken); + + if (_captureMessageContent) + { + SetOutputMessages(activity, response.Messages); + } + } + + return response; + } + catch (Exception ex) + { + if (activity != null) + { + activity.SetStatus(ActivityStatusCode.Error, ex.Message); + activity.AddException(ex); + } + throw; + } + } + + public override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + using var activity = _activitySource.StartActivity("Chat Completion Stream", ActivityKind.Client); + var startTime = DateTime.UtcNow; + bool firstChunkReceived = false; + + if (activity != null) + { + SetSpanType(activity, "llm"); + activity.SetTag("stream", true); + SetModel(activity, options?.ModelId); + + if (_captureMessageContent) + { + SetInputMessages(activity, messages); + } + } + + StringBuilder? outputBuilder = _captureMessageContent ? new() : null; + string? role = null; + string? responseModel = null; + + await using var enumerator = base.GetStreamingResponseAsync(messages, options, cancellationToken) + .GetAsyncEnumerator(cancellationToken); + + while (true) + { + ChatResponseUpdate update; + try + { + if (!await enumerator.MoveNextAsync().ConfigureAwait(false)) + break; + update = enumerator.Current; + } + catch (Exception ex) + { + if (activity != null) + { + activity.SetStatus(ActivityStatusCode.Error, ex.Message); + activity.AddException(ex); + } + throw; + } + + if (!firstChunkReceived && activity != null) + { + SetTimeToFirstToken(activity, (DateTime.UtcNow - startTime).TotalSeconds); + firstChunkReceived = true; + } + + if (outputBuilder != null) + { + role ??= update.Role?.Value; + responseModel ??= update.ModelId; + + if (update.Text != null) + { + outputBuilder.Append(update.Text); + } + } + + yield return update; + } + + if (activity != null && outputBuilder != null) + { + SetResponseModel(activity, responseModel); + var outputJson = ToJson(new object[] + { + new { role = role ?? "assistant", content = outputBuilder.ToString() } + }); + activity.SetTag("braintrust.output_json", outputJson); + } + } + + #region Span Tagging Helpers + + private static void SetSpanType(Activity activity, string spanType) + { + activity.SetTag("braintrust.span_attributes", ToJson(new { type = spanType })); + } + + private static void SetInputMessages(Activity activity, IEnumerable messages) + { + try + { + var input = messages.Select(m => new + { + role = m.Role.Value, + content = m.Text + }); + var json = ToJson(input); + activity.SetTag("braintrust.input_json", json); + } + catch + { + // Ignore serialization errors + } + } + + private static void SetOutputMessages(Activity activity, IList messages) + { + try + { + var output = messages.Select(m => new + { + role = m.Role.Value, + content = m.Text + }); + var json = ToJson(output); + activity.SetTag("braintrust.output_json", json); + } + catch + { + // Ignore serialization errors + } + } + + private static void SetTokenMetrics(Activity activity, UsageDetails? usage) + { + if (usage == null) return; + + if (usage.InputTokenCount.HasValue) + { + activity.SetTag("braintrust.metrics.prompt_tokens", usage.InputTokenCount.Value); + } + if (usage.OutputTokenCount.HasValue) + { + activity.SetTag("braintrust.metrics.completion_tokens", usage.OutputTokenCount.Value); + } + if (usage.TotalTokenCount.HasValue) + activity.SetTag("braintrust.metrics.tokens", usage.TotalTokenCount.Value); + } + + private static void SetTimeToFirstToken(Activity activity, double seconds) + { + if (seconds > 0) + activity.SetTag("braintrust.metrics.time_to_first_token", seconds); + } + + private static void SetModel(Activity activity, string? model) + { + if (model != null) + activity.SetTag("gen_ai.request.model", model); + } + + private static void SetResponseModel(Activity activity, string? model) + { + if (model != null) + activity.SetTag("gen_ai.response.model", model); + } + + private static string? ToJson(T obj) + { + try + { + return JsonSerializer.Serialize(obj, JsonOptions); + } + catch + { + return null; + } + } + + #endregion +} diff --git a/src/Braintrust.Sdk.Extensions.AI/BraintrustExtensionsAI.cs b/src/Braintrust.Sdk.Extensions.AI/BraintrustExtensionsAI.cs new file mode 100644 index 0000000..c3a867d --- /dev/null +++ b/src/Braintrust.Sdk.Extensions.AI/BraintrustExtensionsAI.cs @@ -0,0 +1,103 @@ +using System.Diagnostics; +using Microsoft.Extensions.AI; + +namespace Braintrust.Sdk.Extensions.AI; + +/// +/// Braintrust tracing instrumentation for any IChatClient via Microsoft.Extensions.AI. +/// +/// Provides extension methods to add Braintrust tracing at two pipeline levels: +/// chat client-level (LLM calls) and function-level (tool calls). +/// +public static class BraintrustExtensionsAI +{ + /// + /// Adds Braintrust tracing middleware to a ChatClientBuilder. + /// Creates spans for each LLM call capturing prompts, completions, token usage, and timing. + /// + public static ChatClientBuilder UseBraintrustTracing( + this ChatClientBuilder builder, + bool captureMessageContent = true) + { + var braintrust = Braintrust.Get(); + var activitySource = braintrust.GetActivitySource(); + return builder.UseBraintrustTracing(activitySource, captureMessageContent); + } + + /// + /// Adds Braintrust tracing middleware to a ChatClientBuilder using a custom ActivitySource. + /// + public static ChatClientBuilder UseBraintrustTracing( + this ChatClientBuilder builder, + ActivitySource activitySource, + bool captureMessageContent = true) + { + if (builder == null) + throw new ArgumentNullException(nameof(builder)); + if (activitySource == null) + throw new ArgumentNullException(nameof(activitySource)); + + return builder.Use(innerClient => + new BraintrustChatClient(innerClient, activitySource, captureMessageContent)); + } + + /// + /// Adds both LLM-level and function-level Braintrust tracing to a ChatClientBuilder. + /// + public static ChatClientBuilder UseAllBraintrustTracing( + this ChatClientBuilder builder, + bool captureMessageContent = true, + bool captureToolArguments = true) + { + var braintrust = Braintrust.Get(); + var activitySource = braintrust.GetActivitySource(); + return builder.UseAllBraintrustTracing(activitySource, captureMessageContent, captureToolArguments); + } + + /// + /// Adds both LLM-level and function-level Braintrust tracing using a custom ActivitySource. + /// + public static ChatClientBuilder UseAllBraintrustTracing( + this ChatClientBuilder builder, + ActivitySource activitySource, + bool captureMessageContent = true, + bool captureToolArguments = true) + { + return builder + .UseBraintrustTracing(activitySource, captureMessageContent) + .UseBraintrustFunctionTracing(activitySource, captureToolArguments); + } + + /// + /// Adds Braintrust function call tracing to a ChatClientBuilder. + /// + public static ChatClientBuilder UseBraintrustFunctionTracing( + this ChatClientBuilder builder, + bool captureToolArguments = true) + { + var braintrust = Braintrust.Get(); + var activitySource = braintrust.GetActivitySource(); + return builder.UseBraintrustFunctionTracing(activitySource, captureToolArguments); + } + + /// + /// Adds Braintrust function call tracing using a custom ActivitySource. + /// + public static ChatClientBuilder UseBraintrustFunctionTracing( + this ChatClientBuilder builder, + ActivitySource activitySource, + bool captureToolArguments = true) + { + if (builder == null) + throw new ArgumentNullException(nameof(builder)); + if (activitySource == null) + throw new ArgumentNullException(nameof(activitySource)); + + return builder.UseFunctionInvocation(configure: client => + { + var defaultInvoker = client.FunctionInvoker; + client.FunctionInvoker = BraintrustFunctionMiddleware.CreateInvoker( + activitySource, captureToolArguments, defaultInvoker); + }); + } +} diff --git a/src/Braintrust.Sdk.Extensions.AI/BraintrustFunctionMiddleware.cs b/src/Braintrust.Sdk.Extensions.AI/BraintrustFunctionMiddleware.cs new file mode 100644 index 0000000..acfb421 --- /dev/null +++ b/src/Braintrust.Sdk.Extensions.AI/BraintrustFunctionMiddleware.cs @@ -0,0 +1,119 @@ +using System.Diagnostics; +using System.Text.Json; +using Microsoft.Extensions.AI; +using OpenTelemetry.Trace; + +namespace Braintrust.Sdk.Extensions.AI; + +/// +/// Function calling middleware that wraps tool/function invocations with Braintrust tracing spans. +/// Creates dedicated execute_tool child spans with gen_ai.tool.* attributes — filling a gap +/// that M.E.AI's UseOpenTelemetry() does not cover. +/// +internal static class BraintrustFunctionMiddleware +{ + private static readonly JsonSerializerOptions JsonOptions = new() + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase, + DefaultIgnoreCondition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull + }; + + internal static Func> + CreateInvoker(ActivitySource activitySource, bool captureToolArguments, Func>? defaultInvoker) + { + return async (context, cancellationToken) => + { + var functionName = context.Function?.Name ?? "unknown"; + using var activity = activitySource.StartActivity($"function:{functionName}", ActivityKind.Internal); + var startTime = DateTime.UtcNow; + + try + { + if (activity != null) + { + activity.SetTag("braintrust.span_attributes", ToJson(new { type = "function_call" })); + activity.SetTag("gen_ai.operation.name", "execute_tool"); + activity.SetTag("function.name", functionName); + activity.SetTag("gen_ai.tool.name", functionName); + activity.SetTag("function.iteration", context.Iteration); + activity.SetTag("function.call_index", context.FunctionCallIndex); + activity.SetTag("function.total_count", context.FunctionCount); + + if (captureToolArguments && context.Arguments != null) + { + try + { + var argsJson = ToJson(context.Arguments); + activity.SetTag("braintrust.input_json", argsJson); + activity.SetTag("gen_ai.tool.call.arguments", argsJson); + } + catch + { + // Ignore serialization errors + } + } + } + + object? result; + if (defaultInvoker != null) + { + result = await defaultInvoker(context, cancellationToken).ConfigureAwait(false); + } + else + { + result = context.Function != null + ? await context.Function.InvokeAsync(context.Arguments, cancellationToken).ConfigureAwait(false) + : null; + } + + if (activity != null) + { + var duration = (DateTime.UtcNow - startTime).TotalSeconds; + activity.SetTag("braintrust.metrics.duration", duration); + + if (captureToolArguments && result != null) + { + try + { + var resultJson = ToJson(new { result }); + activity.SetTag("braintrust.output_json", resultJson); + activity.SetTag("gen_ai.tool.call.result", resultJson); + } + catch + { + // Ignore serialization errors + } + } + + if (context.Terminate) + { + activity.SetTag("function.terminated", true); + } + } + + return result; + } + catch (Exception ex) + { + if (activity != null) + { + activity.SetStatus(ActivityStatusCode.Error, ex.Message); + activity.AddException(ex); + } + throw; + } + }; + } + + private static string? ToJson(T obj) + { + try + { + return JsonSerializer.Serialize(obj, JsonOptions); + } + catch + { + return null; + } + } +} diff --git a/src/Braintrust.Sdk.Extensions.AI/README.md b/src/Braintrust.Sdk.Extensions.AI/README.md new file mode 100644 index 0000000..f0f83bf --- /dev/null +++ b/src/Braintrust.Sdk.Extensions.AI/README.md @@ -0,0 +1,16 @@ +# Braintrust.Sdk.Extensions.AI + +Braintrust tracing instrumentation for any `IChatClient` implementation via Microsoft.Extensions.AI. + +## Usage + +```csharp +using Braintrust.Sdk.Extensions.AI; + +var chatClient = new ChatClientBuilder(innerClient) + .UseBraintrustTracing(activitySource) + .UseBraintrustFunctionTracing(activitySource) + .Build(); +``` + +Works with any provider that implements `IChatClient`: OpenAI, Azure OpenAI, Ollama, etc. diff --git a/tests/Braintrust.Sdk.Extensions.AI.Tests/Braintrust.Sdk.Extensions.AI.Tests.csproj b/tests/Braintrust.Sdk.Extensions.AI.Tests/Braintrust.Sdk.Extensions.AI.Tests.csproj new file mode 100644 index 0000000..431b007 --- /dev/null +++ b/tests/Braintrust.Sdk.Extensions.AI.Tests/Braintrust.Sdk.Extensions.AI.Tests.csproj @@ -0,0 +1,24 @@ + + + + net8.0 + enable + enable + false + true + + + + + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + + + + + + diff --git a/tests/Braintrust.Sdk.Extensions.AI.Tests/BraintrustChatClientTests.cs b/tests/Braintrust.Sdk.Extensions.AI.Tests/BraintrustChatClientTests.cs new file mode 100644 index 0000000..c72e91d --- /dev/null +++ b/tests/Braintrust.Sdk.Extensions.AI.Tests/BraintrustChatClientTests.cs @@ -0,0 +1,275 @@ +using System.Diagnostics; +using System.Runtime.CompilerServices; +using Braintrust.Sdk.Extensions.AI; +using Microsoft.Extensions.AI; +using Xunit; + +namespace Braintrust.Sdk.Extensions.AI.Tests; + +public class BraintrustChatClientTests +{ + private static readonly ActivitySource TestSource = new("Braintrust.Tests.ExtensionsAI"); + + [Fact] + public async Task UseBraintrustTracing_CreatesLlmSpan() + { + var activities = new List(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == "Braintrust.Tests.ExtensionsAI", + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => activities.Add(activity) + }; + ActivitySource.AddActivityListener(listener); + + var innerClient = new TestChatClient("LLM response"); + var tracedClient = new ChatClientBuilder(innerClient) + .UseBraintrustTracing(TestSource) + .Build(); + + var messages = new[] { new ChatMessage(ChatRole.User, "Hello") }; + await tracedClient.GetResponseAsync(messages); + + var llmActivity = activities.First(a => a.OperationName == "Chat Completion"); + Assert.Equal(ActivityKind.Client, llmActivity.Kind); + + // Braintrust attributes + var spanType = llmActivity.GetTagItem("braintrust.span_attributes")?.ToString(); + Assert.Contains("\"type\":\"llm\"", spanType); + } + + [Fact] + public async Task UseBraintrustTracing_CapturesInputOutput() + { + var activities = new List(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == "Braintrust.Tests.ExtensionsAI", + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => activities.Add(activity) + }; + ActivitySource.AddActivityListener(listener); + + var innerClient = new TestChatClient("The answer is 42"); + var tracedClient = new ChatClientBuilder(innerClient) + .UseBraintrustTracing(TestSource, captureMessageContent: true) + .Build(); + + var messages = new[] { new ChatMessage(ChatRole.User, "What is the meaning?") }; + await tracedClient.GetResponseAsync(messages); + + var llmActivity = activities.First(a => a.OperationName == "Chat Completion"); + + var inputJson = llmActivity.GetTagItem("braintrust.input_json")?.ToString(); + Assert.NotNull(inputJson); + Assert.Contains("What is the meaning?", inputJson); + + var outputJson = llmActivity.GetTagItem("braintrust.output_json")?.ToString(); + Assert.NotNull(outputJson); + Assert.Contains("The answer is 42", outputJson); + } + + [Fact] + public async Task UseBraintrustTracing_RecordsErrorOnException() + { + var activities = new List(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == "Braintrust.Tests.ExtensionsAI", + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => activities.Add(activity) + }; + ActivitySource.AddActivityListener(listener); + + var innerClient = new TestChatClient(throwException: new HttpRequestException("API error")); + var tracedClient = new ChatClientBuilder(innerClient) + .UseBraintrustTracing(TestSource) + .Build(); + + var messages = new[] { new ChatMessage(ChatRole.User, "Test") }; + await Assert.ThrowsAsync( + () => tracedClient.GetResponseAsync(messages)); + + var llmActivity = activities.First(a => a.OperationName == "Chat Completion"); + Assert.Equal(ActivityStatusCode.Error, llmActivity.Status); + } + + [Fact] + public async Task UseBraintrustTracing_StreamingCreatesSpan() + { + var activities = new List(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == "Braintrust.Tests.ExtensionsAI", + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => activities.Add(activity) + }; + ActivitySource.AddActivityListener(listener); + + var innerClient = new TestChatClient("Streamed response"); + var tracedClient = new ChatClientBuilder(innerClient) + .UseBraintrustTracing(TestSource) + .Build(); + + var messages = new[] { new ChatMessage(ChatRole.User, "Stream test") }; + var updates = new List(); + await foreach (var update in tracedClient.GetStreamingResponseAsync(messages)) + { + updates.Add(update); + } + + Assert.NotEmpty(updates); + var streamActivity = activities.First(a => a.OperationName == "Chat Completion Stream"); + Assert.Equal(ActivityKind.Client, streamActivity.Kind); + Assert.Equal(true, streamActivity.GetTagItem("stream")); + } + + [Fact] + public async Task UseAllBraintrustTracing_CreatesBothLlmAndFunctionSpans() + { + var activities = new List(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == "Braintrust.Tests.ExtensionsAI", + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => activities.Add(activity) + }; + ActivitySource.AddActivityListener(listener); + + var getWeather = AIFunctionFactory.Create((string city) => $"Sunny in {city}", "GetWeather"); + var mockClient = new ToolCallingChatClient(getWeather); + var tracedClient = new ChatClientBuilder(mockClient) + .UseAllBraintrustTracing(TestSource) + .Build(); + + var messages = new List + { + new(ChatRole.User, "What's the weather in Seattle?") + }; + var options = new ChatOptions { Tools = [getWeather] }; + await tracedClient.GetResponseAsync(messages, options); + + // LLM spans + var llmActivities = activities.Where(a => a.OperationName == "Chat Completion").ToList(); + Assert.NotEmpty(llmActivities); + + // Function span with gen_ai attributes + var funcActivity = activities.FirstOrDefault(a => a.OperationName.StartsWith("function:")); + Assert.NotNull(funcActivity); + Assert.Contains("GetWeather", funcActivity.OperationName); + Assert.Equal("execute_tool", funcActivity.GetTagItem("gen_ai.operation.name")); + Assert.Equal("GetWeather", funcActivity.GetTagItem("gen_ai.tool.name")); + } + + [Fact] + public async Task UseBraintrustTracing_SkipsContentWhenDisabled() + { + var activities = new List(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == "Braintrust.Tests.ExtensionsAI", + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => activities.Add(activity) + }; + ActivitySource.AddActivityListener(listener); + + var innerClient = new TestChatClient("Secret response"); + var tracedClient = new ChatClientBuilder(innerClient) + .UseBraintrustTracing(TestSource, captureMessageContent: false) + .Build(); + + var messages = new[] { new ChatMessage(ChatRole.User, "Secret input") }; + await tracedClient.GetResponseAsync(messages); + + var llmActivity = activities.First(a => a.OperationName == "Chat Completion"); + Assert.Null(llmActivity.GetTagItem("braintrust.input_json")); + Assert.Null(llmActivity.GetTagItem("braintrust.output_json")); + } +} + +/// +/// Minimal IChatClient implementation for testing. +/// +internal class TestChatClient : IChatClient +{ + private readonly string _responseText; + private readonly Exception? _exception; + + public TestChatClient(string responseText = "Hello!", Exception? throwException = null) + { + _responseText = responseText; + _exception = throwException; + } + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + if (_exception != null) throw _exception; + + var responseMessage = new ChatMessage(ChatRole.Assistant, _responseText); + return Task.FromResult(new ChatResponse([responseMessage])); + } + + public async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + if (_exception != null) throw _exception; + + yield return new ChatResponseUpdate(ChatRole.Assistant, _responseText); + await Task.CompletedTask; + } + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() { } +} + +/// +/// A chat client that simulates returning a tool call, then a final response. +/// +internal class ToolCallingChatClient : IChatClient +{ + private readonly AIFunction _function; + + public ToolCallingChatClient(AIFunction function) + { + _function = function; + } + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + var messageList = messages.ToList(); + + if (messageList.Any(m => m.Role == ChatRole.Tool)) + { + return Task.FromResult(new ChatResponse( + [new ChatMessage(ChatRole.Assistant, "The weather is sunny.")])); + } + + var functionCallContent = new FunctionCallContent( + callId: "call_1", + name: _function.Name, + arguments: new Dictionary { ["city"] = "Seattle" }); + var assistantMsg = new ChatMessage(ChatRole.Assistant, [functionCallContent]); + return Task.FromResult(new ChatResponse([assistantMsg])); + } + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() { } +}