66import com .microsoft .semantickernel .agents .AgentThread ;
77import com .microsoft .semantickernel .agents .KernelAgent ;
88import com .microsoft .semantickernel .builders .SemanticKernelBuilder ;
9- import com .microsoft .semantickernel .contextvariables .ContextVariable ;
109import com .microsoft .semantickernel .orchestration .InvocationContext ;
1110import com .microsoft .semantickernel .orchestration .InvocationReturnMode ;
1211import com .microsoft .semantickernel .orchestration .PromptExecutionSettings ;
2322import reactor .core .publisher .Flux ;
2423import reactor .core .publisher .Mono ;
2524
25+ import javax .annotation .Nullable ;
2626import java .util .List ;
27- import java .util .Map ;
28- import java .util .function .Function ;
2927import java .util .stream .Collectors ;
3028
3129public class ChatCompletionAgent extends KernelAgent {
3230
33- ChatCompletionAgent (
31+ private ChatCompletionAgent (
3432 String id ,
3533 String name ,
3634 String description ,
@@ -61,43 +59,49 @@ public class ChatCompletionAgent extends KernelAgent {
6159 * @return A Mono containing the agent response
6260 */
6361 @ Override
64- public Mono <List <AgentResponseItem <ChatMessageContent <?>>>> invokeAsync (List <ChatMessageContent <?>> messages , AgentThread thread , AgentInvokeOptions options ) {
65-
66- Mono <ChatHistory > chatHistoryFromThread = this .<ChatHistoryAgentThread >ensureThreadExistsAsync (messages , thread , ChatHistoryAgentThread ::new )
67- .cast (ChatHistoryAgentThread .class )
68- .map (ChatHistoryAgentThread ::getChatHistory )
69- .flatMap (threadChatHistory -> {
70- return Mono .just (new ChatHistory (threadChatHistory .getMessages ()));
71- });
72-
73-
74- Mono <List <ChatMessageContent <?>>> updatedChatHistory = chatHistoryFromThread .flatMap (
75- chatHistory -> internalInvokeAsync (
76- this .getName (),
77- chatHistory ,
78- options
79- )
80- );
62+ public Mono <List <AgentResponseItem <ChatMessageContent <?>>>> invokeAsync (
63+ List <ChatMessageContent <?>> messages ,
64+ AgentThread thread ,
65+ @ Nullable AgentInvokeOptions options
66+ ) {
67+ return ensureThreadExistsWithMessagesAsync (messages , thread , ChatHistoryAgentThread ::new )
68+ .cast (ChatHistoryAgentThread .class )
69+ .flatMap (agentThread -> {
70+ // Extract the chat history from the thread
71+ ChatHistory history = new ChatHistory (
72+ agentThread .getChatHistory ().getMessages ()
73+ );
8174
82- return updatedChatHistory .flatMap (chatMessageContents -> {
83- return Flux .fromIterable (chatMessageContents )
84- .concatMap (chatMessageContent -> this .notifyThreadOfNewMessageAsync (thread , chatMessageContent ))
85- .then (Mono .just (chatMessageContents )); // return the original list
86- }).flatMap (chatMessageContents -> {
87- return Mono .just (chatMessageContents .stream ()
88- .map (chatMessageContent -> {
89- return new AgentResponseItem <ChatMessageContent <?>>(
90- chatMessageContent ,
91- thread );
92- }).collect (Collectors .toList ()));
93- });
75+ // Invoke the agent with the chat history
76+ return internalInvokeAsync (
77+ history ,
78+ options
79+ )
80+ .flatMapMany (Flux ::fromIterable )
81+ // notify on the new thread instance
82+ .concatMap (agentMessage -> {
83+ // Set the author name for the message
84+ agentMessage .setAuthorName (this .name );
85+
86+ return this .notifyThreadOfNewMessageAsync (agentThread , agentMessage ).thenReturn (agentMessage );
87+ })
88+ .collectList ()
89+ .map (chatMessageContents ->
90+ chatMessageContents .stream ()
91+ .map (message -> new AgentResponseItem <ChatMessageContent <?>>(message , agentThread ))
92+ .collect (Collectors .toList ())
93+ );
94+ });
9495 }
9596
9697 private Mono <List <ChatMessageContent <?>>> internalInvokeAsync (
97- String agentName ,
9898 ChatHistory history ,
99- AgentInvokeOptions options
99+ @ Nullable AgentInvokeOptions options
100100 ) {
101+ if (options == null ) {
102+ options = new AgentInvokeOptions ();
103+ }
104+
101105 final Kernel kernel = options .getKernel () != null ? options .getKernel () : this .kernel ;
102106 final KernelArguments arguments = mergeArguments (options .getKernelArguments ());
103107 final String additionalInstructions = options .getAdditionalInstructions ();
@@ -112,13 +116,13 @@ private Mono<List<ChatMessageContent<?>>> internalInvokeAsync(
112116
113117 ToolCallBehavior toolCallBehavior = invocationContext != null
114118 ? invocationContext .getToolCallBehavior ()
115- : ToolCallBehavior .allowAllKernelFunctions (false );
119+ : ToolCallBehavior .allowAllKernelFunctions (true );
116120
117121 // Build base invocation context
118122 InvocationContext .Builder builder = InvocationContext .builder ()
119123 .withPromptExecutionSettings (executionSettings )
120124 .withToolCallBehavior (toolCallBehavior )
121- .withReturnMode (InvocationReturnMode .FULL_HISTORY );
125+ .withReturnMode (InvocationReturnMode .NEW_MESSAGES_ONLY );
122126
123127 if (invocationContext != null ) {
124128 builder = builder
@@ -129,7 +133,7 @@ private Mono<List<ChatMessageContent<?>>> internalInvokeAsync(
129133
130134 InvocationContext agentInvocationContext = builder .build ();
131135
132- return formatInstructionsAsync (kernel , arguments , agentInvocationContext ).flatMap (
136+ return renderInstructionsAsync (kernel , arguments , agentInvocationContext ).flatMap (
133137 instructions -> {
134138 // Create a new chat history with the instructions
135139 ChatHistory chat = new ChatHistory (
@@ -144,24 +148,26 @@ private Mono<List<ChatMessageContent<?>>> internalInvokeAsync(
144148 ));
145149 }
146150
151+ // Add the chat history to the new chat
147152 chat .addAll (history );
148- int previousHistorySize = chat .getMessages ().size ();
149-
150- return chatCompletionService .getChatMessageContentsAsync (chat , kernel , agentInvocationContext )
151- .map (chatMessageContents -> {
152- return chatMessageContents .subList (
153- previousHistorySize ,
154- chatMessageContents .size ());
155- });
153+
154+ return chatCompletionService .getChatMessageContentsAsync (chat , kernel , agentInvocationContext );
156155 }
157156 );
158157
159-
160158 } catch (ServiceNotFoundException e ) {
161159 throw new RuntimeException (e );
162160 }
163161 }
164162
163+
164+ @ Override
165+ public Mono <Void > notifyThreadOfNewMessageAsync (AgentThread thread , ChatMessageContent <?> message ) {
166+ return Mono .defer (() -> {
167+ return thread .onNewMessageAsync (message );
168+ });
169+ }
170+
165171 /**
166172 * Builder for creating instances of ChatCompletionAgent.
167173 */
@@ -174,71 +180,123 @@ public static class Builder implements SemanticKernelBuilder<ChatCompletionAgent
174180 private String name ;
175181 private String description ;
176182 private Kernel kernel ;
177- private KernelArguments KernelArguments ;
183+ private KernelArguments kernelArguments ;
178184 private InvocationContext invocationContext ;
179185 private String instructions ;
180186 private PromptTemplate template ;
181187
188+ /**
189+ * Set the ID of the agent.
190+ *
191+ * @param id The ID of the agent.
192+ */
182193 public Builder withId (String id ) {
183194 this .id = id ;
184195 return this ;
185196 }
186197
198+ /**
199+ * Set the name of the agent.
200+ *
201+ * @param name The name of the agent.
202+ */
187203 public Builder withName (String name ) {
188204 this .name = name ;
189205 return this ;
190206 }
191207
208+ /**
209+ * Set the description of the agent.
210+ *
211+ * @param description The description of the agent.
212+ */
192213 public Builder withDescription (String description ) {
193214 this .description = description ;
194215 return this ;
195216 }
196217
218+ /**
219+ * Set the kernel to use for the agent.
220+ *
221+ * @param kernel The kernel to use.
222+ */
197223 public Builder withKernel (Kernel kernel ) {
198224 this .kernel = kernel ;
199225 return this ;
200226 }
201227
228+ /**
229+ * Set the kernel arguments to use for the agent.
230+ *
231+ * @param KernelArguments The kernel arguments to use.
232+ */
202233 public Builder withKernelArguments (KernelArguments KernelArguments ) {
203- this .KernelArguments = KernelArguments ;
234+ this .kernelArguments = KernelArguments ;
204235 return this ;
205236 }
206237
238+ /**
239+ * Set the instructions for the agent.
240+ *
241+ * @param instructions The instructions for the agent.
242+ */
207243 public Builder withInstructions (String instructions ) {
208244 this .instructions = instructions ;
209245 return this ;
210246 }
211247
248+ /**
249+ * Set the invocation context for the agent.
250+ *
251+ * @param invocationContext The invocation context to use.
252+ */
212253 public Builder withInvocationContext (InvocationContext invocationContext ) {
213254 this .invocationContext = invocationContext ;
214255 return this ;
215256 }
216257
258+ /**
259+ * Set the template for the agent.
260+ *
261+ * @param template The template to use.
262+ */
217263 public Builder withTemplate (PromptTemplate template ) {
218264 this .template = template ;
219265 return this ;
220266 }
221267
268+ /**
269+ * Build the ChatCompletionAgent instance.
270+ *
271+ * @return The ChatCompletionAgent instance.
272+ */
222273 public ChatCompletionAgent build () {
223274 return new ChatCompletionAgent (
224275 id ,
225276 name ,
226277 description ,
227278 kernel ,
228- KernelArguments ,
279+ kernelArguments ,
229280 invocationContext ,
230281 instructions ,
231282 template
232283 );
233284 }
234285
286+ /**
287+ * Build the ChatCompletionAgent instance with the given prompt template config and factory.
288+ *
289+ * @param promptTemplateConfig The prompt template config to use.
290+ * @param promptTemplateFactory The prompt template factory to use.
291+ * @return The ChatCompletionAgent instance.
292+ */
235293 public ChatCompletionAgent build (PromptTemplateConfig promptTemplateConfig , PromptTemplateFactory promptTemplateFactory ) {
236294 return new ChatCompletionAgent (
237295 id ,
238296 name ,
239297 description ,
240298 kernel ,
241- KernelArguments ,
299+ kernelArguments ,
242300 invocationContext ,
243301 promptTemplateConfig .getTemplate (),
244302 promptTemplateFactory .tryCreate (promptTemplateConfig )
0 commit comments