mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Preserve base AgentRunOptions properties and avoid double-wrap on user factory
Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/3afbf18c-de22-4236-a2f2-02ca1e98ae21 Co-authored-by: rogerbarreto <19890735+rogerbarreto@users.noreply.github.com>
This commit is contained in:
committed by
Roger Barreto
Unverified
parent
ca732733d3
commit
921162425d
@@ -229,24 +229,46 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
|
||||
}
|
||||
|
||||
string? sourceName = this._sourceName;
|
||||
IChatClient WrapWithOpenTelemetry(IChatClient cc) =>
|
||||
cc.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build();
|
||||
static IChatClient WrapIfNeeded(IChatClient cc, string? sourceName) =>
|
||||
cc.GetService(typeof(OpenTelemetryChatClient)) is not null
|
||||
? cc
|
||||
: cc.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build();
|
||||
|
||||
if (options is ChatClientAgentRunOptions ccOptions)
|
||||
{
|
||||
// Don't mutate the caller's options; clone and chain any caller-provided factory.
|
||||
// If the user factory already returns an OpenTelemetry-instrumented client, don't double-wrap.
|
||||
var clone = (ChatClientAgentRunOptions)ccOptions.Clone();
|
||||
var userFactory = clone.ChatClientFactory;
|
||||
clone.ChatClientFactory = cc => WrapWithOpenTelemetry(userFactory is null ? cc : userFactory(cc));
|
||||
clone.ChatClientFactory = cc => WrapIfNeeded(userFactory is null ? cc : userFactory(cc), sourceName);
|
||||
return clone;
|
||||
}
|
||||
|
||||
return new ChatClientAgentRunOptions
|
||||
// For a plain AgentRunOptions (or null), create a ChatClientAgentRunOptions and preserve
|
||||
// any base AgentRunOptions properties from the caller so they reach the inner agent.
|
||||
var newOptions = new ChatClientAgentRunOptions
|
||||
{
|
||||
ChatClientFactory = WrapWithOpenTelemetry,
|
||||
ChatClientFactory = cc => WrapIfNeeded(cc, sourceName),
|
||||
};
|
||||
|
||||
if (options is not null)
|
||||
{
|
||||
CopyBaseAgentRunOptions(options, newOptions);
|
||||
}
|
||||
|
||||
return newOptions;
|
||||
}
|
||||
|
||||
#pragma warning disable MEAI001 // ContinuationToken is experimental; copy it through to preserve caller-provided value.
|
||||
private static void CopyBaseAgentRunOptions(AgentRunOptions source, AgentRunOptions target)
|
||||
{
|
||||
target.ContinuationToken = source.ContinuationToken;
|
||||
target.AllowBackgroundResponses = source.AllowBackgroundResponses;
|
||||
target.AdditionalProperties = source.AdditionalProperties?.Clone();
|
||||
target.ResponseFormat = source.ResponseFormat;
|
||||
}
|
||||
#pragma warning restore MEAI001
|
||||
|
||||
/// <summary>The stub <see cref="IChatClient"/> used to delegate from the <see cref="OpenTelemetryChatClient"/> into the inner <see cref="AIAgent"/>.</summary>
|
||||
/// <param name="parentAgent"></param>
|
||||
private sealed class ForwardingChatClient(OpenTelemetryAgent parentAgent) : IChatClient
|
||||
|
||||
@@ -792,6 +792,75 @@ public class OpenTelemetryAgentTests
|
||||
Assert.Contains(activities, a => string.Equals(a.GetTagItem("gen_ai.operation.name") as string, "chat", StringComparison.Ordinal));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task AutoWireChatClient_PlainAgentRunOptions_PreservesBaseProperties_Async()
|
||||
{
|
||||
// Auto-wiring converts a plain AgentRunOptions into a ChatClientAgentRunOptions. The base
|
||||
// properties (ContinuationToken, AllowBackgroundResponses, AdditionalProperties, ResponseFormat)
|
||||
// must be preserved so they reach the inner agent.
|
||||
AgentRunOptions? observedOptions = null;
|
||||
var fakeChatClient = new AutoWireTestChatClient();
|
||||
var innerChatClientAgent = new ChatClientAgent(fakeChatClient);
|
||||
|
||||
// Wrapping agent: surfaces the ChatClientAgent via GetService (so auto-wiring activates),
|
||||
// but captures the AgentRunOptions passed to RunAsync by the OpenTelemetryAgent.
|
||||
var wrapper = new TestAIAgent
|
||||
{
|
||||
GetServiceFunc = (type, key) =>
|
||||
type == typeof(ChatClientAgent) ? innerChatClientAgent : null,
|
||||
RunAsyncFunc = (messages, session, options, ct) =>
|
||||
{
|
||||
observedOptions = options;
|
||||
return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "ok")));
|
||||
},
|
||||
};
|
||||
|
||||
using var agent = new OpenTelemetryAgent(wrapper);
|
||||
|
||||
var additionalProps = new AdditionalPropertiesDictionary { ["customKey"] = "customValue" };
|
||||
var inputOptions = new AgentRunOptions
|
||||
{
|
||||
AllowBackgroundResponses = true,
|
||||
AdditionalProperties = additionalProps,
|
||||
ResponseFormat = ChatResponseFormat.Json,
|
||||
};
|
||||
|
||||
_ = await agent.RunAsync("hi", options: inputOptions);
|
||||
|
||||
Assert.NotNull(observedOptions);
|
||||
Assert.IsType<ChatClientAgentRunOptions>(observedOptions);
|
||||
Assert.Equal(true, observedOptions!.AllowBackgroundResponses);
|
||||
Assert.Same(ChatResponseFormat.Json, observedOptions.ResponseFormat);
|
||||
Assert.NotNull(observedOptions.AdditionalProperties);
|
||||
Assert.Equal("customValue", observedOptions.AdditionalProperties!["customKey"]);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task AutoWireChatClient_UserFactoryReturnsInstrumentedClient_DoesNotDoubleWrap_Async()
|
||||
{
|
||||
var sourceName = Guid.NewGuid().ToString();
|
||||
var activities = new List<Activity>();
|
||||
using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
|
||||
.AddSource(sourceName)
|
||||
.AddInMemoryExporter(activities)
|
||||
.Build();
|
||||
|
||||
var fakeChatClient = new AutoWireTestChatClient();
|
||||
var inner = new ChatClientAgent(fakeChatClient);
|
||||
using var agent = new OpenTelemetryAgent(inner, sourceName);
|
||||
|
||||
// User factory wraps the chat client with OpenTelemetryChatClient itself.
|
||||
var runOptions = new ChatClientAgentRunOptions
|
||||
{
|
||||
ChatClientFactory = cc => cc.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build(),
|
||||
};
|
||||
|
||||
_ = await agent.RunAsync("hi", options: runOptions);
|
||||
|
||||
// Expect 2 activities (invoke_agent + a single chat span). If we double-wrapped, we would see 3.
|
||||
Assert.Equal(2, activities.Count);
|
||||
}
|
||||
|
||||
private sealed class AutoWireTestChatClient : IChatClient
|
||||
{
|
||||
public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default) =>
|
||||
|
||||
Reference in New Issue
Block a user