AG-UI support in Agent Framework #4347

jiapingzeng · 2025-10-24T23:56:53Z

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:

Enable streaming:

POST /_cluster/settings
{
    "persistent": {
        "plugins.ml_commons.stream_enabled": true,
        "plugins.ml_commons.ag_ui_enabled": true
    }
}

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:-}}"
            }
        ]
    }
}

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:

Feature flag disabled:

data: {"error": "Error processing request: The AG-UI agent feature is not enabled. To enable, please update the setting plugins.ml_commons.ag_ui_enabled"}

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"}

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"}

.../main/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandler.java

pyek-bot · 2025-11-19T06:13:26Z

common/src/main/java/org/opensearch/ml/common/agui/RunErrorEvent.java

+        builder.field("message", message);
+        if (code != null) {
+            builder.field("code", code);
+        }


can message be null?

pyek-bot · 2025-11-19T06:15:16Z

common/src/main/java/org/opensearch/ml/common/agui/RunErrorEvent.java

+
+    public RunErrorEvent(String message, String code) {
+        super(TYPE, System.currentTimeMillis(), null);
+        this.message = message != null ? message : "";


similar comment as below, can we just keep this as null? instead of empty and then not write or read if null

pyek-bot · 2025-11-19T06:17:11Z

common/src/main/java/org/opensearch/ml/common/agui/RunErrorEvent.java

+@EqualsAndHashCode(callSuper = true)
+public class RunErrorEvent extends BaseEvent {
+
+    public static final String TYPE = "RUN_ERROR";


nit: can be enum

common/src/main/java/org/opensearch/ml/common/agui/TextMessageStartEvent.java

pyek-bot · 2025-11-19T06:24:04Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AGUIFrontendTool.java

+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {


you can throw an exception instead of error message

pyek-bot · 2025-11-19T06:26:05Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java

+            wrappedTools.put(toolName, frontendToolObj);
+        }
+
+        return wrappedTools;


This might be missing a bracket }

good catch, not sure how that happened

pyek-bot · 2025-11-19T06:28:13Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAGUIAgentRunner.java

+                            } else {
+                                log.error("AG-UI: Invalid function calling configuration: {}", llmInterface);
+                            }


should this fail in this case? what happens here?

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAGUIAgentRunner.java

pyek-bot · 2025-11-19T06:52:28Z

.../main/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandler.java


+            // Initialize AG-UI message state for this LLM response
+            if (isAGUIAgent) {
+                messageId = "msg_" + System.currentTimeMillis() + "_" + System.nanoTime();


can just use nano tiime? is this some standard?

plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteAction.java

plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteStreamAction.java

.../main/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandler.java

pyek-bot · 2025-11-19T07:07:21Z

.../src/main/java/org/opensearch/ml/engine/function_calling/BedrockConverseFunctionCalling.java

+
+                assistantMessage.getContent().add(Map.of("toolUse", toolUse));
+            }
+        } catch (Exception e) {


log the error?

pyek-bot · 2025-11-19T07:08:54Z

plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteStreamAction.java

+            String threadId = inputDataSet.getParameters().get(AGUI_PARAM_THREAD_ID);
+            return threadId != null ? threadId : "thread_" + System.currentTimeMillis();
+        }
+        return "thread_" + System.currentTimeMillis();


will it cause a problem if the IDs are same in case multiple calls at made at the same time?

no, IDs are just dummy values at the moment. When we integrate this with memory we can start using proper IDs.

Signed-off-by: Jiaping Zeng <[email protected]>

jiapingzeng · 2025-11-19T09:41:46Z

thanks, will address other comments in follow up PR

jiapingzeng changed the title ~~AG-UI support~~ AG-UI support in Agent Framework Oct 24, 2025

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 24, 2025 23:59 — with GitHub Actions Error

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 24, 2025 23:59 — with GitHub Actions Failure

jiapingzeng force-pushed the ag-ui-support branch from 05dfe19 to f555403 Compare October 25, 2025 00:00

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Error

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Failure

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Error

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:01 — with GitHub Actions Failure

jiapingzeng commented Oct 25, 2025

View reviewed changes

.../main/java/org/opensearch/ml/engine/algorithms/remote/streaming/BedrockStreamingHandler.java Outdated Show resolved Hide resolved

jiapingzeng force-pushed the ag-ui-support branch from f555403 to 4721dc5 Compare October 25, 2025 00:26

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Failure

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Error

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 00:27 — with GitHub Actions Failure

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Error

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Failure

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Error

jiapingzeng had a problem deploying to ml-commons-cicd-env-require-approval October 25, 2025 09:38 — with GitHub Actions Failure

jiapingzeng requested a deployment to ml-commons-cicd-env-require-approval October 27, 2025 10:23 — with GitHub Actions Waiting

jiapingzeng force-pushed the ag-ui-support branch from 013b69d to f8f6d12 Compare October 27, 2025 17:18

jiapingzeng marked this pull request as ready for review October 27, 2025 17:18

jiapingzeng requested review from b4sjoo, dhrubo-os, jngz-es and mingshl as code owners October 27, 2025 17:18