Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
54 changes: 52 additions & 2 deletions dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,7 @@ public async Task<ChatResponse> GetResponseAsync(
var response = await parentAgent.InnerAgent.RunAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false);

// Wrap the response in a ChatResponse so we can pass it back through OpenTelemetryChatClient.
return response.AsChatResponse();
return ToChatResponse(response);
}

public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
Expand All @@ -312,10 +312,60 @@ public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false))
{
// Wrap the response updates in ChatResponseUpdates so we can pass them back through OpenTelemetryChatClient.
yield return update.AsChatResponseUpdate();
yield return ToChatResponseUpdate(update);
}
}

/// <summary>
/// Wraps an <see cref="AgentResponse"/> in a <see cref="ChatResponse"/> that keeps the original response as its
/// <see cref="ChatResponse.RawRepresentation"/>, so <see cref="RunCoreAsync"/> can hand the
/// inner agent's own response back to the caller unchanged.
/// </summary>
/// <remarks>
/// <see cref="AgentResponseExtensions.AsChatResponse"/> isn't used here because it returns the inner agent's
/// <see cref="ChatResponse"/> directly whenever that is the response's raw representation. Round-tripping
/// through that <see cref="ChatResponse"/> discards agent-level state such as
/// <see cref="AgentResponse.ContinuationToken"/> and <see cref="AgentResponse.AgentId"/>, which for a
/// <see cref="ChatClientAgent"/> means the caller receives the underlying provider's continuation token rather
/// than the agent's own, and resuming with it fails.
/// </remarks>
private static ChatResponse ToChatResponse(AgentResponse response) =>
new()
{
AdditionalProperties = response.AdditionalProperties,
ContinuationToken = response.ContinuationToken,
ConversationId = (response.RawRepresentation as ChatResponse)?.ConversationId,
Comment thread
atty57 marked this conversation as resolved.
Resolved
CreatedAt = response.CreatedAt,
FinishReason = response.FinishReason,
Messages = response.Messages,
RawRepresentation = response,
ResponseId = response.ResponseId,
Usage = response.Usage,
};

/// <summary>
/// Wraps an <see cref="AgentResponseUpdate"/> in a <see cref="ChatResponseUpdate"/> that keeps the original
/// update as its <see cref="ChatResponseUpdate.RawRepresentation"/>, so
/// <see cref="RunCoreStreamingAsync"/> can hand the inner agent's own update back to the
/// caller unchanged. See <see cref="ToChatResponse"/> for why
/// <see cref="AgentResponseExtensions.AsChatResponseUpdate"/> isn't used.
/// </summary>
private static ChatResponseUpdate ToChatResponseUpdate(AgentResponseUpdate update) =>
new()
{
AdditionalProperties = update.AdditionalProperties,
AuthorName = update.AuthorName,
Contents = update.Contents,
ContinuationToken = update.ContinuationToken,
ConversationId = (update.RawRepresentation as ChatResponseUpdate)?.ConversationId,
CreatedAt = update.CreatedAt,
FinishReason = update.FinishReason,
MessageId = update.MessageId,
RawRepresentation = update,
ResponseId = update.ResponseId,
Role = update.Role,
Comment thread
SergeyMenshykh marked this conversation as resolved.
Resolved
};

public object? GetService(Type serviceType, object? serviceKey = null) =>
// Delegate any inquiries made by the OpenTelemetryChatClient back to the parent agent.
parentAgent.GetService(serviceType, serviceKey);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -945,6 +945,81 @@ public async Task AutoWireChatClient_PlainAgentRunOptions_PreservesContinuationT
Assert.Same(inputOptions, observedOptions);
Assert.Same(token, observedOptions.ContinuationToken);
}

[Fact]
public async Task RunAsync_WrappingChatClientAgent_PreservesAgentContinuationTokenAsync()
{
// Arrange
// The inner ChatClientAgent wraps the chat client's token in a ChatClientAgentContinuationToken. That
// wrapper must survive the AgentResponse -> ChatResponse -> AgentResponse round trip that this agent
// performs through OpenTelemetryChatClient, otherwise the caller gets the provider's raw token back and
// resuming with it fails.
var innerToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
var fakeChatClient = new AutoWireTestChatClient { ContinuationToken = innerToken };
var innerAgent = new ChatClientAgent(fakeChatClient);
using var agent = new OpenTelemetryAgent(innerAgent);
var session = await agent.CreateSessionAsync();

// Act
var response = await agent.RunAsync("hi", session, new AgentRunOptions { AllowBackgroundResponses = true });

// Assert
var agentToken = Assert.IsType<ChatClientAgentContinuationToken>(response.ContinuationToken);
Assert.Same(innerToken, agentToken.InnerToken);
Assert.Equal(innerAgent.Id, response.AgentId);
}

[Fact]
public async Task RunAsync_WrappingChatClientAgent_ResumesFromSerializedContinuationTokenAsync()
{
// Arrange
var innerToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
ChatOptions? observedChatOptions = null;
var fakeChatClient = new AutoWireTestChatClient
{
ContinuationToken = innerToken,
OnGetResponseAsync = (_, opts) => observedChatOptions = opts,
};
using var agent = new OpenTelemetryAgent(new ChatClientAgent(fakeChatClient));
var session = await agent.CreateSessionAsync();

var first = await agent.RunAsync("hi", session, new AgentRunOptions { AllowBackgroundResponses = true });

// Act
// Serialize and deserialize the token the way a caller resuming the run from another process would.
Assert.NotNull(first.ContinuationToken);
var restoredToken = ResponseContinuationToken.FromBytes(first.ContinuationToken.ToBytes());
_ = await agent.RunAsync(session, new AgentRunOptions { ContinuationToken = restoredToken });

// Assert
// The agent must have unwrapped the restored token and passed the chat client's own token back down.
Assert.NotNull(observedChatOptions?.ContinuationToken);
Assert.Equal(innerToken.ToBytes().ToArray(), observedChatOptions.ContinuationToken.ToBytes().ToArray());
}

[Fact]
public async Task RunStreamingAsync_WrappingChatClientAgent_PreservesAgentContinuationTokenAsync()
{
// Arrange
var innerToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 });
var fakeChatClient = new AutoWireTestChatClient { ContinuationToken = innerToken };
var innerAgent = new ChatClientAgent(fakeChatClient);
using var agent = new OpenTelemetryAgent(innerAgent);
var session = await agent.CreateSessionAsync();

// Act
var updates = new List<AgentResponseUpdate>();
await foreach (var update in agent.RunStreamingAsync("hi", session, new AgentRunOptions { AllowBackgroundResponses = true }))
{
updates.Add(update);
}

// Assert
var tokenUpdate = Assert.Single(updates, u => u.ContinuationToken is not null);
var agentToken = Assert.IsType<ChatClientAgentContinuationToken>(tokenUpdate.ContinuationToken);
Assert.Same(innerToken, agentToken.InnerToken);
Assert.Equal(innerAgent.Id, tokenUpdate.AgentId);
}
#pragma warning restore MEAI001

[Fact]
Expand Down Expand Up @@ -1198,18 +1273,23 @@ private sealed class AutoWireTestChatClient : IChatClient
{
public Action<IEnumerable<ChatMessage>, ChatOptions?>? OnGetResponseAsync { get; set; }

#pragma warning disable MEAI001 // ResponseContinuationToken is experimental.
/// <summary>Gets or sets the continuation token to report on the response, simulating a background response.</summary>
public ResponseContinuationToken? ContinuationToken { get; set; }

public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
{
this.OnGetResponseAsync?.Invoke(messages, options);
return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")));
return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")) { ContinuationToken = this.ContinuationToken });
}

public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
this.OnGetResponseAsync?.Invoke(messages, options);
await Task.Yield();
yield return new ChatResponseUpdate(ChatRole.Assistant, "ok");
yield return new ChatResponseUpdate(ChatRole.Assistant, "ok") { ContinuationToken = this.ContinuationToken };
}
#pragma warning restore MEAI001

public object? GetService(Type serviceType, object? serviceKey = null) =>
serviceType?.IsInstanceOfType(this) == true ? this : null;
Expand Down
Loading