-
Notifications
You must be signed in to change notification settings - Fork 186
AG-UI support in Agent Framework #4347
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
05dfe19 to
f555403
Compare
|
Still working on clean up and adding UTs but features are ready. |
| ? StaticCredentialsProvider | ||
| .create(AwsSessionCredentials.create(connector.getAccessKey(), connector.getSecretKey(), connector.getSessionToken())) | ||
| : StaticCredentialsProvider.create(AwsBasicCredentials.create(connector.getAccessKey(), connector.getSecretKey())); | ||
| return java.security.AccessController.doPrivileged((java.security.PrivilegedAction<BedrockRuntimeAsyncClient>) () -> { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
was getting this error locally:
java.lang.SecurityException: Denied OPEN (read) access to file: /Users/jpz/.aws/credentials, domain: ProtectionDomain (file:/Volumes/workplace/github/ml-commons/plugin/build/testclusters/integTest-0/distro/3.4.0-ARCHIVE/lib/opensearch-3.4.0-SNAPSHOT.jar <no signer certificates>)
Not sure why the Bedrock client was trying to access my .credentials when credentials are already provided via connector, will remove this doPrivileged after testing in cluster to see if I run into the same issue.
f555403 to
4721dc5
Compare
013b69d to
f8f6d12
Compare
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
Signed-off-by: Jiaping Zeng <[email protected]>
0ead4f9 to
db7a8cb
Compare
|
|
||
| @Override | ||
| protected void addEventSpecificFields(XContentBuilder builder, Params params) throws IOException { | ||
| builder.field("messageId", messageId); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[nit] let's add these as a constant variable.
| } | ||
| } | ||
|
|
||
| private void processAGUIMessages(MLAgent mlAgent, Map<String, String> params, String llmInterface) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This method is soo long, can we break this down for simplicity?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
May be something like:
• extractToolResults(JsonArray messageArray)
• buildChatHistory(JsonArray messageArray, Map<String, String> params)
• processMessageTemplates(JsonArray messageArray, Map<String, String> params)
| frontendTools.addAll(parsed); | ||
| } | ||
| } catch (Exception e) { | ||
| log.warn("Failed to parse frontend tools: {}", e.getMessage()); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
what about log.error?
| } | ||
|
|
||
| @Override | ||
| public boolean validate(Map<String, String> parameters) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we add a todo here if we are planning to add validation later?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We don't expect this tool to be executed here. It will be executed in browser which handles validation. This class is to let the LLM knows that it has access to frontend tools.
| || !jsonObj.has(AGUI_FIELD_RUN_ID) | ||
| || !jsonObj.has(AGUI_FIELD_MESSAGES) | ||
| || !jsonObj.has(AGUI_FIELD_TOOLS)) { | ||
| return false; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we add a log?
| public TextMessageContentEvent(String messageId, String delta) { | ||
| super(TYPE, System.currentTimeMillis(), null); | ||
| this.messageId = messageId; | ||
| this.delta = delta != null ? delta : ""; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why do we need have "". Can't we just mark this as optional string?
|
|
||
| @Override | ||
| public String getType() { | ||
| return "AGUIFrontendTool"; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's put all these in static variable.
| import lombok.extern.log4j.Log4j2; | ||
|
|
||
| @Log4j2 | ||
| public class AGUIInputConverter { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't see any comprehensive input sanitization. This can be potential for injection attacks through malformed AG-UI JSON input
| List<Map<String, Object>> frontendTools = parseFrontendTools(aguiTools); | ||
|
|
||
| // Process with unified tools (both frontend and backend) | ||
| processUnifiedTools(mlAgent, params, listener, memory, sessionId, functionCalling, frontendTools); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
what did you mean by rocess with unified tools (both frontend and backend), where are we sending backend tools? Also is there any chance functionality of front end tool and backend tool can be overlapped?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Frontend tools need separate processing as we need to break out of ReAct loop and send response to frontend. However, the LLM needs to know that it has both frontend and backend tools for it to generate a response. This logic passes all tools into the LLM.
I created an RFC here that goes over the flow for frontend and backend tools: #4409
| ); | ||
| // Check if this is a backend tool - if it is, execute it normally in the ReAct loop | ||
| // If it's NOT a backend tool, it must be a frontend tool, so break out of the loop | ||
| boolean isBackendTool = backendTools != null && backendTools.containsKey(action); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
So frontend tool and back end tool can't co-exist together?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
explained in above thread
Description
Add AG-UI support in Agent Framework
Related Issues
Resolves #4409
Check List
--signoff.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.
Sample setup:
Sample requests:
request:
response:
request:
response: