Skip to content

Commit 4c331c1

Browse files
alliscodeBen Thomas
andauthored
.Net: Bedrock dotnet (#11166)
Improving integration tests for Bedrock Agents ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [ ] The code builds clean without any errors or warnings - [ ] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [ ] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄 --------- Co-authored-by: Ben Thomas <[email protected]>
1 parent 52cbc33 commit 4c331c1

File tree

10 files changed

+220
-119
lines changed

10 files changed

+220
-119
lines changed

dotnet/samples/GettingStartedWithAgents/BedrockAgent/Step01_BedrockAgent.cs

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

3+
using Microsoft.SemanticKernel;
4+
using Microsoft.SemanticKernel.Agents;
35
using Microsoft.SemanticKernel.Agents.Bedrock;
6+
using Microsoft.SemanticKernel.ChatCompletion;
47

58
namespace GettingStarted.BedrockAgents;
69

@@ -22,17 +25,19 @@ public async Task UseNewAgentAsync()
2225
var bedrockAgent = await this.CreateAgentAsync("Step01_BedrockAgent");
2326

2427
// Respond to user input
28+
AgentThread bedrockAgentThread = new BedrockAgentThread(this.RuntimeClient);
2529
try
2630
{
27-
var responses = bedrockAgent.InvokeAsync(BedrockAgent.CreateSessionId(), UserQuery, null);
28-
await foreach (var response in responses)
31+
var responses = bedrockAgent.InvokeAsync(new ChatMessageContent(AuthorRole.User, UserQuery), bedrockAgentThread, null);
32+
await foreach (ChatMessageContent response in responses)
2933
{
3034
this.Output.WriteLine(response.Content);
3135
}
3236
}
3337
finally
3438
{
3539
await bedrockAgent.Client.DeleteAgentAsync(new() { AgentId = bedrockAgent.Id });
40+
await bedrockAgentThread.DeleteAsync();
3641
}
3742
}
3843

@@ -50,10 +55,18 @@ public async Task UseExistingAgentAsync()
5055
var bedrockAgent = new BedrockAgent(getAgentResponse.Agent, this.Client, this.RuntimeClient);
5156

5257
// Respond to user input
53-
var responses = bedrockAgent.InvokeAsync(BedrockAgent.CreateSessionId(), UserQuery, null);
54-
await foreach (var response in responses)
58+
AgentThread bedrockAgentThread = new BedrockAgentThread(this.RuntimeClient);
59+
try
60+
{
61+
var responses = bedrockAgent.InvokeAsync(new ChatMessageContent(AuthorRole.User, UserQuery), bedrockAgentThread, null);
62+
await foreach (ChatMessageContent response in responses)
63+
{
64+
this.Output.WriteLine(response.Content);
65+
}
66+
}
67+
finally
5568
{
56-
this.Output.WriteLine(response.Content);
69+
await bedrockAgentThread.DeleteAsync();
5770
}
5871
}
5972

@@ -66,19 +79,21 @@ public async Task UseNewAgentStreamingAsync()
6679
{
6780
// Create the agent
6881
var bedrockAgent = await this.CreateAgentAsync("Step01_BedrockAgent_Streaming");
82+
AgentThread bedrockAgentThread = new BedrockAgentThread(this.RuntimeClient);
6983

7084
// Respond to user input
7185
try
7286
{
73-
var streamingResponses = bedrockAgent.InvokeStreamingAsync(BedrockAgent.CreateSessionId(), UserQuery, null);
74-
await foreach (var response in streamingResponses)
87+
var streamingResponses = bedrockAgent.InvokeStreamingAsync(new ChatMessageContent(AuthorRole.User, UserQuery), bedrockAgentThread, null);
88+
await foreach (StreamingChatMessageContent response in streamingResponses)
7589
{
7690
this.Output.WriteLine(response.Content);
7791
}
7892
}
7993
finally
8094
{
8195
await bedrockAgent.Client.DeleteAgentAsync(new() { AgentId = bedrockAgent.Id });
96+
await bedrockAgentThread.DeleteAsync();
8297
}
8398
}
8499

dotnet/samples/GettingStartedWithAgents/BedrockAgent/Step02_BedrockAgent_CodeInterpreter.cs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22

33
using System.Reflection;
44
using Microsoft.SemanticKernel;
5+
using Microsoft.SemanticKernel.Agents;
56
using Microsoft.SemanticKernel.Agents.Bedrock;
7+
using Microsoft.SemanticKernel.ChatCompletion;
68

79
namespace GettingStarted.BedrockAgents;
810

@@ -28,13 +30,14 @@ public async Task UseAgentWithCodeInterpreterAsync()
2830
{
2931
// Create the agent
3032
var bedrockAgent = await this.CreateAgentAsync("Step02_BedrockAgent_CodeInterpreter");
33+
AgentThread bedrockAgentThread = new BedrockAgentThread(this.RuntimeClient);
3134

3235
// Respond to user input
3336
try
3437
{
3538
BinaryContent? binaryContent = null;
36-
var responses = bedrockAgent.InvokeAsync(BedrockAgent.CreateSessionId(), UserQuery, null);
37-
await foreach (var response in responses)
39+
var responses = bedrockAgent.InvokeAsync(new ChatMessageContent(AuthorRole.User, UserQuery), bedrockAgentThread, null);
40+
await foreach (ChatMessageContent response in responses)
3841
{
3942
if (response.Content != null)
4043
{
@@ -71,6 +74,7 @@ public async Task UseAgentWithCodeInterpreterAsync()
7174
finally
7275
{
7376
await bedrockAgent.Client.DeleteAgentAsync(new() { AgentId = bedrockAgent.Id });
77+
await bedrockAgentThread.DeleteAsync();
7478
}
7579
}
7680

dotnet/samples/GettingStartedWithAgents/BedrockAgent/Step03_BedrockAgent_Functions.cs

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
using System.ComponentModel;
44
using Microsoft.SemanticKernel;
55
using Microsoft.SemanticKernel.Agents.Bedrock;
6+
using Microsoft.SemanticKernel.ChatCompletion;
67

78
namespace GettingStarted.BedrockAgents;
89

@@ -25,10 +26,9 @@ public async Task UseAgentWithFunctionsAsync()
2526
try
2627
{
2728
var responses = bedrockAgent.InvokeAsync(
28-
BedrockAgent.CreateSessionId(),
29-
"What is the weather in Seattle?",
29+
new ChatMessageContent(AuthorRole.User, "What is the weather in Seattle?"),
3030
null);
31-
await foreach (var response in responses)
31+
await foreach (ChatMessageContent response in responses)
3232
{
3333
if (response.Content != null)
3434
{
@@ -57,10 +57,9 @@ public async Task UseAgentWithFunctionsComplexTypeAsync()
5757
try
5858
{
5959
var responses = bedrockAgent.InvokeAsync(
60-
BedrockAgent.CreateSessionId(),
61-
"What is the special soup and how much does it cost?",
60+
new ChatMessageContent(AuthorRole.User, "What is the special soup and how much does it cost?"),
6261
null);
63-
await foreach (var response in responses)
62+
await foreach (ChatMessageContent response in responses)
6463
{
6564
if (response.Content != null)
6665
{
@@ -88,10 +87,9 @@ public async Task UseAgentStreamingWithFunctionsAsync()
8887
try
8988
{
9089
var streamingResponses = bedrockAgent.InvokeStreamingAsync(
91-
BedrockAgent.CreateSessionId(),
92-
"What is the weather forecast in Seattle?",
90+
new ChatMessageContent(AuthorRole.User, "What is the weather forecast in Seattle?"),
9391
null);
94-
await foreach (var response in streamingResponses)
92+
await foreach (StreamingChatMessageContent response in streamingResponses)
9593
{
9694
if (response.Content != null)
9795
{
@@ -119,10 +117,9 @@ public async Task UseAgentWithParallelFunctionsAsync()
119117
try
120118
{
121119
var responses = bedrockAgent.InvokeAsync(
122-
BedrockAgent.CreateSessionId(),
123-
"What is the current weather in Seattle and what is the weather forecast in Seattle?",
120+
new ChatMessageContent(AuthorRole.User, "What is the current weather in Seattle and what is the weather forecast in Seattle?"),
124121
null);
125-
await foreach (var response in responses)
122+
await foreach (ChatMessageContent response in responses)
126123
{
127124
if (response.Content != null)
128125
{

dotnet/src/Agents/Bedrock/BedrockAgent.cs

Lines changed: 28 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -75,13 +75,22 @@ public static string CreateSessionId()
7575

7676
#region public methods
7777

78-
// TODO: Add overload that allows the AgentAliasId to be specified
78+
/// <inheritdoc/>
79+
public override IAsyncEnumerable<AgentResponseItem<ChatMessageContent>> InvokeAsync(
80+
ICollection<ChatMessageContent> messages,
81+
AgentThread? thread = null,
82+
AgentInvokeOptions? options = null,
83+
CancellationToken cancellationToken = default)
84+
{
85+
return this.InvokeAsync(messages, thread, options, WorkingDraftAgentAlias, cancellationToken);
86+
}
7987

8088
/// <inheritdoc/>
81-
public override async IAsyncEnumerable<AgentResponseItem<ChatMessageContent>> InvokeAsync(
89+
public async IAsyncEnumerable<AgentResponseItem<ChatMessageContent>> InvokeAsync(
8290
ICollection<ChatMessageContent> messages,
8391
AgentThread? thread = null,
8492
AgentInvokeOptions? options = null,
93+
string? agentAliasId = null,
8594
[EnumeratorCancellation] CancellationToken cancellationToken = default)
8695
{
8796
Verify.NotNull(messages, nameof(messages));
@@ -105,7 +114,7 @@ public override async IAsyncEnumerable<AgentResponseItem<ChatMessageContent>> In
105114

106115
var invokeAgentRequest = new InvokeAgentRequest
107116
{
108-
AgentAliasId = WorkingDraftAgentAlias,
117+
AgentAliasId = agentAliasId ?? WorkingDraftAgentAlias,
109118
SessionState = sessionState,
110119
AgentId = this.Id,
111120
SessionId = bedrockThread.Id,
@@ -118,6 +127,7 @@ public override async IAsyncEnumerable<AgentResponseItem<ChatMessageContent>> In
118127
// Return the results to the caller in AgentResponseItems.
119128
await foreach (var result in invokeResults.ConfigureAwait(false))
120129
{
130+
await this.NotifyThreadOfNewMessage(bedrockThread, result, cancellationToken).ConfigureAwait(false);
121131
yield return new(result, bedrockThread);
122132
}
123133
}
@@ -238,16 +248,21 @@ async IAsyncEnumerable<ChatMessageContent> InvokeInternal()
238248
innerContents.Add(message.InnerContent);
239249
}
240250

241-
yield return content.Length == 0
242-
? throw new KernelException("No content was returned from the agent.")
243-
: new ChatMessageContent(AuthorRole.Assistant, content)
244-
{
245-
AuthorName = this.GetDisplayName(),
246-
Items = items,
247-
ModelId = this.AgentModel.FoundationModel,
248-
Metadata = metadata,
249-
InnerContent = innerContents,
250-
};
251+
if (content.Length == 0)
252+
{
253+
throw new KernelException("No content was returned from the agent.");
254+
}
255+
256+
var chatMessageContent = new ChatMessageContent(AuthorRole.Assistant, content)
257+
{
258+
AuthorName = this.GetDisplayName(),
259+
Items = items,
260+
ModelId = this.AgentModel.FoundationModel,
261+
Metadata = metadata,
262+
InnerContent = innerContents,
263+
};
264+
265+
yield return chatMessageContent;
251266
}
252267
}
253268

@@ -331,22 +346,6 @@ async IAsyncEnumerable<StreamingChatMessageContent> InvokeInternal()
331346

332347
#endregion
333348

334-
/// <inheritdoc/>
335-
protected override Task<TThreadType> EnsureThreadExistsWithMessagesAsync<TThreadType>(ICollection<ChatMessageContent> messages, AgentThread? thread, Func<TThreadType> constructThread, CancellationToken cancellationToken)
336-
{
337-
if (thread is null)
338-
{
339-
thread = constructThread();
340-
}
341-
342-
if (thread is not TThreadType concreteThreadType)
343-
{
344-
throw new KernelException($"{this.GetType().Name} currently only supports agent threads of type {nameof(TThreadType)}.");
345-
}
346-
347-
return Task.FromResult(concreteThreadType);
348-
}
349-
350349
/// <inheritdoc/>
351350
protected override IEnumerable<string> GetChannelKeys()
352351
{
Lines changed: 15 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,9 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

33
using System;
4-
using System.Collections.Generic;
5-
using System.Diagnostics.CodeAnalysis;
6-
using System.Linq;
7-
using System.Runtime.CompilerServices;
84
using System.Threading;
95
using System.Threading.Tasks;
106
using Amazon.BedrockAgentRuntime;
11-
using Amazon.BedrockAgentRuntime.Model;
127

138
namespace Microsoft.SemanticKernel.Agents.Bedrock;
149
/// <summary>
@@ -37,47 +32,7 @@ public BedrockAgentThread(AmazonBedrockAgentRuntimeClient runtimeClient, string?
3732
/// <returns>A task that completes when the thread has been created.</returns>
3833
public new Task CreateAsync(CancellationToken cancellationToken = default)
3934
{
40-
return this.CreateInternalAsync(cancellationToken);
41-
}
42-
43-
/// <summary>
44-
/// Asynchronously retrieves all messages in the thread.
45-
/// </summary>
46-
/// <param name="maxResults">The maximum number of results to return in the response.</param>
47-
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
48-
/// <returns>The messages in the thread.</returns>
49-
/// <exception cref="InvalidOperationException">The thread has been deleted.</exception>
50-
[Experimental("SKEXP0110")]
51-
public async IAsyncEnumerable<ChatMessageContent> GetMessagesAsync(int maxResults = 100, [EnumeratorCancellation] CancellationToken cancellationToken = default)
52-
{
53-
if (this.IsDeleted)
54-
{
55-
throw new InvalidOperationException("This thread has been deleted and cannot be used anymore.");
56-
}
57-
58-
if (this.Id is null)
59-
{
60-
await this.CreateAsync(cancellationToken).ConfigureAwait(false);
61-
}
62-
63-
var invocationSteps = await this._runtimeClient.ListInvocationStepsAsync(new ListInvocationStepsRequest() { SessionIdentifier = this.Id, MaxResults = maxResults }, cancellationToken).ConfigureAwait(false);
64-
var invocationStepTasks = invocationSteps.InvocationStepSummaries.Select(s => this._runtimeClient.GetInvocationStepAsync(new GetInvocationStepRequest { InvocationIdentifier = s.InvocationId }));
65-
await Task.WhenAll(invocationStepTasks).ConfigureAwait(false);
66-
67-
foreach (var invocationStep in invocationStepTasks)
68-
{
69-
var response = await invocationStep.ConfigureAwait(false);
70-
if (response.InvocationStep?.Payload is not null)
71-
{
72-
foreach (BedrockSessionContentBlock? block in response.InvocationStep.Payload.ContentBlocks)
73-
{
74-
yield return new ChatMessageContent
75-
{
76-
Content = block.Text
77-
};
78-
}
79-
}
80-
}
35+
return base.CreateAsync(cancellationToken);
8136
}
8237

8338
/// <inheritdoc />
@@ -90,6 +45,7 @@ public async IAsyncEnumerable<ChatMessageContent> GetMessagesAsync(int maxResult
9045
var response = await this._runtimeClient.CreateSessionAsync(
9146
request: new(),
9247
cancellationToken: cancellationToken).ConfigureAwait(false);
48+
9349
return response.SessionId;
9450
}
9551
catch (AmazonBedrockAgentRuntimeException ex)
@@ -109,12 +65,21 @@ protected override async Task DeleteInternalAsync(CancellationToken cancellation
10965

11066
try
11167
{
112-
var response = await this._runtimeClient.DeleteSessionAsync(
68+
var endSessionResponse = await this._runtimeClient.EndSessionAsync(
69+
request: new()
70+
{
71+
SessionIdentifier = this.Id
72+
},
73+
cancellationToken: cancellationToken).ConfigureAwait(false);
74+
75+
var deleteSessionResponse = await this._runtimeClient.DeleteSessionAsync(
11376
request: new()
11477
{
11578
SessionIdentifier = this.Id
11679
},
11780
cancellationToken: cancellationToken).ConfigureAwait(false);
81+
82+
this.Id = null;
11883
}
11984
catch (AmazonBedrockAgentRuntimeException ex)
12085
{
@@ -127,9 +92,9 @@ protected override async Task DeleteInternalAsync(CancellationToken cancellation
12792
}
12893

12994
/// <inheritdoc />
130-
protected override Task OnNewMessageInternalAsync(ChatMessageContent newMessage, CancellationToken cancellationToken = default)
95+
protected override async Task OnNewMessageInternalAsync(ChatMessageContent newMessage, CancellationToken cancellationToken = default)
13196
{
132-
string message = $"{nameof(BedrockAgentThread)} does not support adding messages directly to the thread.";
133-
throw new NotImplementedException(message);
97+
// Create the thread if it does not exist. Bedrock agents cannot add messages to the thread without invoking so we don't do that here
98+
await this.CreateAsync(cancellationToken).ConfigureAwait(false);
13499
}
135100
}

0 commit comments

Comments
 (0)