AG-UI support in Agent Framework

[jiapingzeng]

Description

Add AG-UI support in Agent Framework

Related Issues

Resolves #4211

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

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

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

3. 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]

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

[jiapingzeng]

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.