Giving AI Agents the power of ML Models with Amazon SageMaker AI and Amazon Bedrock AgentCore

Learn how to use ML models in AI Agents using Amazon Bedrock AgentCore Gateway/Runtime, Model Context Protocol (MCP) and Amazon SageMaker AI endpoints.

Davide Gallitelli (AWS)

Published Aug 4, 2025

Last modified Aug 6, 2025

Code is available on GitHub - ML Models as Tools with AgentCore

In this blog post, you will learn how to leverage a ML model for demand forecasting with your AI Agent. We will use:

Amazon SageMaker AI endpoints to deploy the ML model (XGBoost)
Amazon Bedrock AgentCore Gateway and Model Context Protocol (MCP) to interact with the endpoint using AWS API Smithy Models
Amazon Bedrock AgentCore Runtime and Model Context Protocol (MCP) to interact with the endpoint using custom MCP tools

architecture-diagram — Architecture Diagram for leveraging ML models with Amazon Bedrock AgentCore

Setup - Train and Deploy your model

Before you can get started building your ML-enhanced AI Agent, you will need to train and deploy a Machine Learning model using Amazon SageMaker AI. We've done that plenty of times already, so we'll skip step-by-step instructions here. What you need to know is:

Use a SageMaker Training job to train your model
Deploy your model on a SageMaker AI endpoint (real-time)

The code for this is available in the repository . In our example, we train a demand forecasting model using XGBoost, but any other use case is equally valid.

Next Step - Connect the endpoint to an Agent

Now comes to fun part. You're going to build your AI Agent as a business expert that can leverage the models you've trained to generate accurate predictions. The AI Agent is powered by your LLM of choice, provided by your model provider of choice, as long as it can perform tool calling. If you don't know what tool calling is, I've explained it in this Medium blog . In a nutshell, you let the LLM decide which action to take, perform it, and evaluate its result, using a standard interface. Model Context Protocol (MCP) has set itself as the industry standard for this.

So, our endpoint (and the model on it) will be the tool we need to invoke to generate the prediction. With Amazon AgentCore, we have two options to achieve this:

use Amazon Bedrock AgentCore Gateway as proxy between AWS APIs and our agent:
1. we can expose the Amazon SageMaker AI InvokeEndpoint API as tool to be called via MCP
2. this is easily achievable, since AWS has published API Models in Smithy format, and AgentCore Gateway can natively create MCP Tools starting from the Smithy model definition
3. it's generic enough that it doesn't require us to create a custom MCP server for each endpoint, however will require us some context engineering to make sure the right endpoint is called with the right input
use Amazon Bedrock AgentCore Runtime to host a custom MCP server that interacts with our endpoint:
1. Runtime will take care of the managed deployment of the MCP server
2. We need to build our MCP server which can be annoying if you don't know how to write that code
3. Build the custom MCP server also means that the tool will be super specific to our endpoint, with all the description and inputs which simplify the task of choosing the right tool for the LLM

We will go through both examples in this blog. You can choose whichever you think makes more sense for your use case and your experience.

Amazon SageMaker AI and Amazon Bedrock AgentCore Gateway

Code is available here: GitHub

Bedrock AgentCore Gateway provides customers a way to turn their existing Smithy APIs into fully-managed MCP servers without needing to manage infra or hosting. Customers can bring Smithy spec and transform them into mcp tools. We will demonstrate creating mcp tools from smithy model of Amazon S3. The agent will then be able to query Amazon S3 and answer questions related to the it.

The Gateway workflow involves the following steps to connect your agents to external tools:

Create the tools for your Gateway - Define your tools using Smithy specification.
Create a Gateway endpoint - Create the gateway that will serve as the MCP entry point with inbound authentication.
Add targets to your Gateway - Configure the Smithy target that define how the gateway routes requests to specific tools. All the operations that part of Smithy file will become an MCP-compatible tool, and will be made available through your Gateway endpoint URL.
Update your agent code - Connect your agent to the Gateway endpoint to access all configured tools through the unified MCP interface.

Since AWS API Models are available on GitHub , we can focus on how to turn SageMaker Runtime Smithy API Model into a MCP tool:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import boto3

# Create the agentcore client
agentcore_client = boto3.client('bedrock-agentcore-control')

# Create a Smithy model target
target = agentcore_client.create_gateway_target(
    gatewayIdentifier=GATEWAY_ID,
    name="SageMakerRuntimeTarget",
    targetConfiguration={
        "mcp": {
            "smithyModel": {
                "inlinePayload": open('sagemaker-runtime-2017-05-13.json', 'r').read()
            }
        }
    },
    credentialProviderConfigurations=[
        {
            "credentialProviderType": "GATEWAY_IAM_ROLE"
        }
    ]
)

Then, you can build your custom AI Agent using your framework of choice, in our case Strands Agents:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
mcp_client = MCPClient(lambda: create_streamable_http_transport(mcp_url, access_token))
     
with mcp_client:
    tools = get_full_tools_list(mcp_client)
    print(f"Found the following tools: {[tool.tool_name for tool in tools]}")

    # Create an agent with these tools
    agent = Agent(
        model="us.anthropic.claude-3-5-haiku-20241022-v1:0", tools=tools,
        system_prompt=f"""
            You are an expert data analyst and a product owner. Your task is to help the business thrive.

            # Sytems access
            You have access to a ML model for demand forecasting, exposed by the Amazon SageMaker AI endpoint {ENDPOINT_NAME}.
            This models takes an array of float as input, and returns an array of float as output.
        """
    )
    
    # Invoke the agent
    agent(
        f"These are the current demand values:\n\n<input>{test_sample}</input>\n\n"
        "Predict demand, please. Provide the output in JSON format {'predictions':<predictions>}"
    )

Notice how the system prompt has to be written in such a way that the agent knows about the existence of a ml-models-as-tools endpoint, and that it can be used for demand forecasting. If you have a catalog with a list of endpoints and tasks that it can be used for, this system prompt can be skipped by providing an additional tool to access such catalog.

Amazon SageMaker AI and Amazon Bedrock AgentCore Runtime

Code is available here: GitHub

In this section of the blog, you will learn how to host MCP (Model Context Protocol) servers on Amazon Bedrock AgentCore Runtime. We will use the Amazon Bedrock AgentCore Python SDK to wrap the code to interact with the Amazon SageMaker AI endpoint as an MCP server compatible with Amazon Bedrock AgentCore. The Amazon Bedrock AgentCore Python SDK handles the MCP server implementation details so you can focus on your tools' core functionality. It transforms your code into the AgentCore standardized MCP protocol contracts for direct communication.The first step is to write the MCP Server code, and test it locally provided you have the right IAM permissions. For this implementation, we decided to use FastMCP to define the MCP server, Streamable HTTP as the transport protocol, and boto3 to connect to the SageMaker AI endpoint.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
from mcp.server.fastmcp import FastMCP
from starlette.responses import JSONResponse
import json
import boto3
import numpy as np

mcp = FastMCP(host="0.0.0.0", stateless_http=True)

@mcp.tool()
def demand_forecasting_with_ml(sample_data):
	"""
	This tool predicts future demand based on historical data using a machine learning model.
	"""
	sagemaker_runtime = boto3.client("sagemaker-runtime")
	endpoint_name = "ml-models-as-tools"
	response = sagemaker_runtime.invoke_endpoint(
		EndpointName=endpoint_name,
		Body=json.dumps(sample_data),
		ContentType="application/json",
		Accept="application/json"
	)
	predictions = json.loads(response['Body'].read().decode("utf-8"))
	return np.array(predictions)

if __name__ == "__main__":
	mcp.run(transport="streamable-http")

You can test this MCP server locally by saving the code in a mcp_server.py file, and running python mcp_server.py. Then, use MCP Inspector or an AI agent written in your preferred agentic framework to connect to the MCP server and list the tools. Once your tests are complete, you now need to deploy the MCP server to a production-ready service, like Amazon Bedrock AgentCore Runtime:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from bedrock_agentcore_starter_toolkit import Runtime

agentcore_runtime = Runtime()
response = agentcore_runtime.configure(
    agent_name=tool_name,
    region=region,
    auto_create_ecr=True,   # Instructs Runtime to auto-create ECR repo
    entrypoint="mcp_server.py", # The core of the MCP server
    requirements_file="requirements.txt", # dependencies
    execution_role="IAM-ROLE-WITH-PERMISSIONS-TO-INVOKE-SAGEMAKER-AI",
    authorizer_configuration={
        "customJWTAuthorizer": {
            "allowedClients": [LIST-OF-ALLOWED-CLIENTS],
            "discoveryUrl": "DISCOVERY-URL",
        }
    },
    protocol="MCP", # Coming soon: support for A2A
)

launch_result = agentcore_runtime.launch()
print("Launch completed ✓")
print(f"Agent ARN: {launch_result.agent_arn}")
print(f"Agent ID: {launch_result.agent_id}")

After a couple of minutes, your MCP server will be available for queries! Before we can invoke our deployed MCP server, let's create the MCP URL of the server, which will require the ARN of the resource we just created (confusingly called Agent ARN):

1
2
3
server_runtime_arn = launch_result.agent_arn
encoded_arn = server_runtime_arn.replace(':', '%3A').replace('/', '%2F')
mcp_url = f"https://bedrock-agentcore.{region}.amazonaws.com/runtimes/{encoded_arn}/invocations?qualifier=DEFAULT"

Now, you can use this MCP URL with your MCP-compatible Agentic AI framework of choice, in my case Strands Agents:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
from strands import Agent
from strands.tools.mcp.mcp_client import MCPClient
from utils.cognito_utils import get_token
from utils.agent_utils import create_streamable_http_transport,get_full_tools_list 
import os

def run_agent(mcp_url: str, access_token: str):
    mcp_client = MCPClient(lambda: create_streamable_http_transport(mcp_url, access_token))
     
    with mcp_client:
        tools = get_full_tools_list(mcp_client)
        print(f"Found the following tools: {[tool.tool_name for tool in tools]}")

        # Create an agent with these tools
        agent = Agent(
            model="us.anthropic.claude-3-5-haiku-20241022-v1:0", tools=tools
        )
        
        # Invoke the agent
        agent(
            f"These are the current demand values:\n\n<input>{test_sample}</input>\n\n"
            "Predict demand, please. Provide the output in JSON format {'predictions':<predictions>}."
            "Only reply with the prediction JSON, nothing else. If the tool fails, just tell me the error."
        )
        
run_agent(<MCP URL>, <Access token>)

As you can see with respect to the previous example with AgentCore Gateway, now the MCP server already already knows and exposes the detail that the tool will allow the agent to interact with a demand forecasting algorithm. This means that we don't need to provide this information at runtime as part of a system prompt, and instead it's dynamically discovered and added to the context when the tools are retrieved from the MCP server itself.

Thanks for checking out this blog! If you made it here, you have learnt how to:

✅️ Deploy a demand forecasting model (XGBoost) to Amazon SageMaker AI endpoints
✅️ Use MCP to connect to an Amazon SageMaker AI endpoint
✅️ Use AgentCore Gateway with AWS API Smithy Models to have a scalable, secure proxy server to your resources (in our case, an Amazon SageMaker AI endpoint)
✅️ Use AgentCore Runtime to host an MCP server with a custom tool to invoke the ML model hosted on the SageMaker AI endpoint

# agentic-ai # amazon-sagemaker-ai # generative-ai # machine-learning # amazon-bedrock-agentcore

Any opinions in this article are those of the individual author and may not reflect the opinions of AWS.

Enjoyed reading this content? Let the author know!

Your likes, comments, shares, and saves help creators reach more builders.

Loading recommendations

Loading article

Builder Center

Giving AI Agents the power of ML Models with Amazon SageMaker AI and Amazon Bedrock AgentCore

Davide Gallitelli (AWS)

Setup - Train and Deploy your model

Next Step - Connect the endpoint to an Agent

Amazon SageMaker AI and Amazon Bedrock AgentCore Gateway

Amazon SageMaker AI and Amazon Bedrock AgentCore Runtime