Skip to content

Conversation

@jiapingzeng
Copy link
Contributor

@jiapingzeng jiapingzeng commented Oct 24, 2025

Description

Add AG-UI support in Agent Framework

Related Issues

Resolves #4409

Check List

  • New functionality includes testing.
  • New functionality has been documented.
  • API changes companion pull request created.
  • Commits are signed per the DCO using --signoff.
  • Public documentation issue/PR created.

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:

  1. Enable streaming:
POST /_cluster/settings
{
    "persistent": {
        "plugins.ml_commons.stream_enabled": true
    }
}
  1. Model:
{
    "name": "Claude 3.7",
    "function_name": "remote",
    "description": "Bedrock Claude model",
    "connector": {
        "name": "Bedrock Converse Connector",
        "description": "Bedrock Converse Connector",
        "version": 1,
        "protocol": "aws_sigv4",
        "parameters": {
            "region": "us-east-1",
            "model": "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
            "service_name": "bedrock"
        },
        "credential": {
            "access_key": "{{aws-access-key-id}}",
            "secret_key": "{{aws-secret-access-key}}",
            "session_token": "{{aws-session-token}}"
        },
        "actions": [
            {
                "action_type": "predict",
                "method": "POST",
                "url": "https://bedrock-runtime.${parameters.region}.amazonaws.com/model/${parameters.model}/converse-stream",
                "request_body": "{\"messages\": [${parameters._chat_history:-}{\"role\":\"user\",\"content\":[{\"text\":\"${parameters.prompt}\"}]}${parameters._interactions:-}]${parameters.tool_configs:-}}"
            }
        ]
    }
}
  1. Agent:
POST /_plugins/_ml/agents/_register
{
    "name": "AG-UI chat agent",
    "type": "AG_UI",
    "description": "this is a test agent",
    "llm": {
        "model_id": "{{model_id}}",
        "parameters": {
            "max_iteration": 50,
            "system_prompt": "You are a helpful assistant. Use the available tools to help users. When you need to perform an action, use the appropriate tool by calling it with the correct parameters.",
            "prompt": "Context:${parameters.context}\nQuestion:${parameters.question}"
        }
    },
    "memory": {
        "type": "conversation_index"
    },
    "parameters": {
        "_llm_interface": "bedrock/converse/claude"
    },
    "tools": [
        {
            "type": "ListIndexTool",
            "name": "ListIndexTool"
        }
    ]
}

Sample requests:

  1. Using OSD context and frontend tool: [AI] Add experimental AI Chat and Context Provider plugins for OpenSearch Dashboards OpenSearch-Dashboards#10600
    request:
POST /_plugins/_ml/agents/{{agent_id}}/_execute/stream
{
    "threadId": "thread-1760484425669-pj44w6ive",
    "runId": "run-1760484493143-gu67vx1jl",
    "messages": [
        {
            "id": "msg-1760484430561-epujp4bi6",
            "role": "user",
            "content": "execute ppl query 'test'"
        }
    ],
    "tools": [
        {
            "name": "execute_ppl_query",
            "description": "Update the query bar with a PPL query and optionally execute it",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The PPL query to set in the query bar"
                    },
                    "autoExecute": {
                        "type": "boolean",
                        "description": "Whether to automatically execute the query (default: true)"
                    },
                    "description": {
                        "type": "string",
                        "description": "Optional description of what the query does"
                    }
                },
                "required": [
                    "query"
                ]
            }
        }
    ],
    "context": [
        {
            "description": "Explore application page context",
            "value": "{\"appId\":\"explore\",\"timeRange\":{\"from\":\"now-15m\",\"to\":\"now\"},\"query\":{\"query\":\"\",\"language\":\"PPL\"}}"
        }
    ],
    "state": {},
    "forwardedProps": {}
}

response:

data: {"type":"RUN_STARTED","timestamp":1761437522976,"threadId":"jxnbHZoBSmgeTrPdltR1","runId":"kBnbHZoBSmgeTrPdltTP"}

data: {"type":"TEXT_MESSAGE_START","timestamp":1761437522976,"messageId":"msg_1761437522975","role":"assistant"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522976,"messageId":"msg_1761437522975","delta":"I"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522977,"messageId":"msg_1761437522975","delta":"'ll"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522977,"messageId":"msg_1761437522975","delta":" help"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522978,"messageId":"msg_1761437522975","delta":" you execute"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522978,"messageId":"msg_1761437522975","delta":" the"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522979,"messageId":"msg_1761437522975","delta":" PPL query "}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437522979,"messageId":"msg_1761437522975","delta":"'test'."}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437523114,"messageId":"msg_1761437522975","delta":" Let"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437523116,"messageId":"msg_1761437522975","delta":" me"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437523117,"messageId":"msg_1761437522975","delta":" "}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437523157,"messageId":"msg_1761437522975","delta":"do that for you right"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437523202,"messageId":"msg_1761437522975","delta":" away"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437523234,"messageId":"msg_1761437522975","delta":"."}

data: {"type":"TEXT_MESSAGE_END","timestamp":1761437524255,"messageId":"msg_1761437522975"}

data: {"type":"TOOL_CALL_START","timestamp":1761437524255,"toolCallId":"tooluse_cDQH5-zFR8SEa9eljI5Xpg","toolCallName":"execute_ppl_query"}

data: {"type":"TOOL_CALL_ARGS","timestamp":1761437524255,"toolCallId":"tooluse_cDQH5-zFR8SEa9eljI5Xpg","delta":"{\"query\":\"test\",\"autoExecute\":true}"}

data: {"type":"TOOL_CALL_END","timestamp":1761437524255,"toolCallId":"tooluse_cDQH5-zFR8SEa9eljI5Xpg"}

data: {"type":"RUN_FINISHED","timestamp":1761437524256,"threadId":"jxnbHZoBSmgeTrPdltR1","runId":"kBnbHZoBSmgeTrPdltTP"}

  1. Frontend sending tool result back:
    request:
POST /_plugins/_ml/agents/{{agent_id}}/_execute/stream
{
    "threadId": "thread-1761299262323-66ekfq52m",
    "runId": "run-1761299273064-qjycuyoo5",
    "messages": [
        {
            "id": "msg-1761299270344-wfxawmjgw",
            "role": "user",
            "content": "execute ppl 'test'"
        },
        {
            "id": "msg_1761299271843",
            "role": "assistant",
            "toolCalls": [
                {
                    "id": "tooluse_gdZkro87QqW01isjNiDEqg",
                    "type": "function",
                    "function": {
                        "name": "execute_ppl_query",
                        "arguments": "{\"query\":\"test\"}"
                    }
                }
            ],
            "content": "I'll help you execute a PPL query with the term 'test'. Let me do that for you."
        },
        {
            "id": "msg-1761299273064-6sg6hlkop",
            "role": "tool",
            "content": "{\"success\":true,\"executed\":true,\"query\":\"test\",\"message\":\"Query updated and executed\"}",
            "toolCallId": "tooluse_gdZkro87QqW01isjNiDEqg"
        }
    ],
    "tools": [
        {
            "name": "execute_ppl_query",
            "description": "Update the query bar with a PPL query and optionally execute it",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The PPL query to set in the query bar"
                    },
                    "autoExecute": {
                        "type": "boolean",
                        "description": "Whether to automatically execute the query (default: true)"
                    },
                    "description": {
                        "type": "string",
                        "description": "Optional description of what the query does"
                    }
                },
                "required": [
                    "query"
                ]
            }
        }
    ],
    "context": [
        {
            "description": "Explore application page context",
            "value": "{\"appId\":\"explore\",\"timeRange\":{\"from\":\"now-15m\",\"to\":\"now\"},\"query\":{\"query\":\"test\",\"language\":\"PPL\"}}"
        }
    ],
    "state": {},
    "forwardedProps": {}
}

response:

data: {"type":"RUN_STARTED","timestamp":1761437671115,"threadId":"kRndHZoBSmgeTrPd2dRj","runId":"khndHZoBSmgeTrPd2dSt"}

data: {"type":"TEXT_MESSAGE_START","timestamp":1761437671116,"messageId":"msg_1761437671115","role":"assistant"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671116,"messageId":"msg_1761437671115","delta":"I"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671173,"messageId":"msg_1761437671115","delta":"'ve"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671189,"messageId":"msg_1761437671115","delta":" executed the PPL"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671320,"messageId":"msg_1761437671115","delta":" query "}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671322,"messageId":"msg_1761437671115","delta":"'test' for"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671323,"messageId":"msg_1761437671115","delta":" you. The"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671406,"messageId":"msg_1761437671115","delta":" query"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671409,"messageId":"msg_1761437671115","delta":" has"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671493,"messageId":"msg_1761437671115","delta":" been successfully"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671494,"messageId":"msg_1761437671115","delta":" submitte"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671527,"messageId":"msg_1761437671115","delta":"d an"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671615,"messageId":"msg_1761437671115","delta":"d execute"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671617,"messageId":"msg_1761437671115","delta":"d. You"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671661,"messageId":"msg_1761437671115","delta":" shoul"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671697,"messageId":"msg_1761437671115","delta":"d now"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671843,"messageId":"msg_1761437671115","delta":" see the results"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671844,"messageId":"msg_1761437671115","delta":" displaye"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671845,"messageId":"msg_1761437671115","delta":"d in the"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437671870,"messageId":"msg_1761437671115","delta":" Explore application interface"}

data: {"type":"TEXT_MESSAGE_CONTENT","timestamp":1761437672102,"messageId":"msg_1761437671115","delta":"."}

data: {"type":"TEXT_MESSAGE_END","timestamp":1761437672159,"messageId":"msg_1761437671115"}

data: {"type":"RUN_FINISHED","timestamp":1761437672159,"threadId":"kRndHZoBSmgeTrPd2dRj","runId":"khndHZoBSmgeTrPd2dSt"}

@jiapingzeng jiapingzeng changed the title AG-UI support AG-UI support in Agent Framework Oct 24, 2025
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 24, 2025 23:59 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 24, 2025 23:59 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 24, 2025 23:59 — with GitHub Actions Failure
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 24, 2025 23:59 — with GitHub Actions Failure
@jiapingzeng
Copy link
Contributor Author

Still working on clean up and adding UTs but features are ready.

@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Failure
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Failure
? 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>) () -> {
Copy link
Contributor Author

@jiapingzeng jiapingzeng Oct 25, 2025

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.

@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Failure
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Failure
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 27, 2025 10:23 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 27, 2025 10:23 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 27, 2025 10:23 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 27, 2025 10:23 — with GitHub Actions Waiting
@jiapingzeng jiapingzeng marked this pull request as ready for review October 27, 2025 17:18
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]>
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:48 — with GitHub Actions Failure
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:48 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:48 — with GitHub Actions Failure
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:48 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:50 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:50 — with GitHub Actions Failure
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:50 — with GitHub Actions Error
@jiapingzeng jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval November 9, 2025 03:50 — with GitHub Actions Failure

@Override
protected void addEventSpecificFields(XContentBuilder builder, Params params) throws IOException {
builder.field("messageId", messageId);
Copy link
Collaborator

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) {
Copy link
Collaborator

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?

Copy link
Collaborator

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());
Copy link
Collaborator

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) {
Copy link
Collaborator

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?

Copy link
Contributor Author

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;
Copy link
Collaborator

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 : "";
Copy link
Collaborator

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";
Copy link
Collaborator

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 {
Copy link
Collaborator

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);
Copy link
Collaborator

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?

Copy link
Contributor Author

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);
Copy link
Collaborator

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

explained in above thread

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[RFC] AG-UI Support in Agent Framework

2 participants