DEV Community

Cover image for Agent conversation context with AWS Bedrock AgentCore Memory
AWS Community Builders profile image szymon-szym
szymon-szym for AWS Community Builders

Posted on

Agent conversation context with AWS Bedrock AgentCore Memory

#aws #agents #cdk

Photo by Photoholgic on Unsplash

You can seamlessly hook up the conversation context in the stateless agentic environment with AWS Bedrock AgentCore Memory

Goal

I want to create a customer support agent (how original), that is capable of holding the conversation context, and can recall main facts about the user in for example few days later.

# What I am building

For testing purposes, I don't need to build the UI, but the idea is that the user can interact with the agent by sending an HTTP request. Authentication is done using the OAuth flow.
The agent understands the context of the request and reaches out to the AgentCore Memory service. It uses the fetched information to enrich the context of the conversation.

The whole code is available IN THIS REPOSITORY

Implementation

I am using AWS AgentCore Runtime to host the agent, and AWS AgentCore Memory to store facts about users. For the user authentication, I have a Cognito user pool with a configured app client.

When building this POC, I reached out for agentcore starter toolkit, which is a perfect tool for the job.

I use Strands as my agentic framework, mostly because it simplifies integration with the AgentCore environment, e.g., managing memory hooks automatically.

Cognito

I create a Cognito user pool for managing user interactions with the agent. Any provider might be used for this, as we will need only the discovery URI and client ID. I define infrastructure with AWS CDK

import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import * as cognito from "aws-cdk-lib/aws-cognito";

export class AwsStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const userPool = new cognito.UserPool(this, "AgentCoreRuntimeUserPool", {
      selfSignUpEnabled: true,
      autoVerify: { email: true },
      signInAliases: { email: true },
      standardAttributes: {
        email: { required: true, mutable: true },
      },
    });

    const domainPrefix = "test-ac-pool";

    const domain = new cognito.UserPoolDomain(
      this,
      "AgentCoreRuntimeUserPoolDomain",
      {
        userPool,
        cognitoDomain: {
          domainPrefix,
        },
      },
    );

    const userPoolClient = new cognito.UserPoolClient(
      this,
      "AgentCoreRuntimeUserPoolClient",
      {
        userPool,
        generateSecret: true,
        authFlows: {
          userPassword: true,
          adminUserPassword: true,
          userSrp: true,
        },
        oAuth: {
          flows: {
            authorizationCodeGrant: true,
            implicitCodeGrant: true,
          },
          scopes: [
            cognito.OAuthScope.EMAIL,
            cognito.OAuthScope.OPENID,
            cognito.OAuthScope.PROFILE,
          ],
          callbackUrls: ["https://example.com/callback"],
          logoutUrls: ["https://example.com/logout"],
        },
      },
    );

    const discovery = ` https://cognito-idp.us-east-1.amazonaws.com/${userPool.userPoolId}/.well-known/openid-configuration`;

    new cdk.CfnOutput(this, "Discovery URI", {
      value: discovery,
      description: "The discovery URI for the user pool",
    });

    new cdk.CfnOutput(this, "Client ID", {
      value: userPoolClient.userPoolClientId,
      description: "The ID of the user pool client",
    });
  }
}
Enter fullscreen mode Exit fullscreen mode

For now, callback URLs are dummy, as I won't connect the UI yet.

After running cdk deploy, I get the discovery URL and client ID

Agent

I create new directory and run uv init. AgentCore toolkit works with uv out of the box.

I add dependencies needed using uv add <package>
My pyproject.toml looks like this:

[project]
name = "bank-support-agent"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "bedrock-agentcore>=1.0.5",
    "bedrock-agentcore-starter-toolkit>=0.1.28",
    "boto3>=1.40.64",
    "python-dotenv>=1.2.1",
    "strands-agents>=1.14.0",
    "strands-agents-tools>=0.2.13",
]
Enter fullscreen mode Exit fullscreen mode

I start by defining some constants and initializing BedrockAgentCoreApp

# main.py
# -- imports --
REGION = "us-east-1"
MODEL_ID = "amazon.nova-pro-v1:0"
MEMORY_ID = os.getenv("BEDROCK_AGENTCORE_MEMORY_ID")

app = BedrockAgentCoreApp()

session = Session(region_name=REGION)

model = BedrockModel(model_id=MODEL_ID, boto_session=session)

# ...
Enter fullscreen mode Exit fullscreen mode

For simplicity's sake, I will keep everything in the one function. Moving forward, I would extract at least initialization part. AgentCore comes with a handy decorator to define the entry point

# ...

@app.entrypoint
def invoke(payload, context):

# ...
Enter fullscreen mode Exit fullscreen mode

I extract actor_id session_id and prompt from the request payload and context.

# ...
    actor_id = headers.get(
        "x-amzn-bedrock-agentcore-runtime-custom-actorid", "default_actor"
    )

    session_id = getattr(context, "session_id", None)

    prompt = payload.get("prompt", "")
# ...
Enter fullscreen mode Exit fullscreen mode

Now I create the Memory config. I opted for a default setup, but AgentCore Memory is highly configurable. We might extend built-in strategies with overrides or use our custom strategy. This is useful, e.g., if we need to collect only data about a specific business topic.

The most important properties for managing memories are session_id and actor_id. Session ID is responsible for the short-term memory, for understanding the current conversation context. The AgentCore Runtime is stateless, so it won't be possible without some storage mechanism. Actor ID is responsible for identifying the caller - user, or other agent, and for getting long-term memories related to the given persona. In our case, we will use built-in strategies for storing facts and preferences.

As I mentioned before, the Strands framework comes with ready-to-use integrations:

# ...

session_manager = None
    if MEMORY_ID:
        session_manager = AgentCoreMemoryConfig(
            memory_id=MEMORY_ID,
            session_id=session_id or "default",
            actor_id=actor_id,
            retrieval_config={
                f"/users/{actor_id}/facts": RetrievalConfig(
                    top_k=3, relevance_score=0.5
                ),
                f"/users/{actor_id}/preferences": RetrievalConfig(
                    top_k=3, relevance_score=0.5
                ),
            },
        )
        session_manager = AgentCoreMemorySessionManager(session_manager, REGION)

# ...
Enter fullscreen mode Exit fullscreen mode

The next step would be to initiate the agent and to respond to the prompt:

 # ...
agent = Agent(
        model=model,
        session_manager=session_manager,
        system_prompt="""
        You are a helpful assistant. Your task is to support bank clients with their queries and concerns.
        Stay cohesive and return short answers.
    """,
        tools=[],
    )

    result = agent(prompt)

    return {"response": result.message.get("content", [{}])[0].get("text", str(result))}
Enter fullscreen mode Exit fullscreen mode

As you see, it's a pretty basic agent definition.

And at the end, let's trigger the app

# ...

if __name__ == "__main__":
    app.run()
Enter fullscreen mode Exit fullscreen mode

Prepare AgentCore config

Now it is time for the agentcore toolkit magic.

In the agent directory, I run uv run agentcore configure -e main.py -r us-east-1

Tool properly detects pyproject.toml. I let it create execturion role and ECR repository.

When it comes to Auth configuration, I enter my discovery URL ...

.. and client ID. Leaving audience empty

I allow X-Amzn-Bedrock-AgentCore-Runtime-Custom-ActorId and Authorization headers to be passed with the request. Finally I choose "yes" for creating memory.

Now the .bedrock_agentcore.yaml is create in root agent directory, and .bedrock_agentcore/bank_customer_service/Dockerfile is ready to be deployed.

Provide cloud resources

I run

uv run agentcore launch
Enter fullscreen mode Exit fullscreen mode

The tool takes care of creating all necessare resources. It also triggers CodeBuild job to build the agent and push the image to ECR. Once it is built, it will create agent in the runtime, or will update the existing one with the new version

Once deployment is ready, I can see, that .bedrock_agentcore.yaml is updated with information about created cloud resources

Testing

There is a convenient method to test agents, using invoke command, e.g.

agentcore invoke '{"prompt": "Hello"}'
Enter fullscreen mode Exit fullscreen mode

However, it won't work for us:

Getting Cognito token

This is not the most exciting part, but until there is no UI to log in, it might be problematic. To be able to get a token, I perform a few actions.

I need a secret hash to be used for creating, confirming, and getting a token using aws cli

echo -n "test_user@example.com<client_id>" | openssl dgst -sha256 -hmac <client_secret> -binary | openssl enc -base64
Enter fullscreen mode Exit fullscreen mode

Now I can create a new user and confirm creation

aws cognito-idp sign-up --client-id <client_id> --username test_user@example.com --password Test1234! --region us-east-1 --secret-hash <hash>
aws cognito-idp admin-confirm-sign-up --user-pool-id <pool_id> --username test_user@example.com --region us-east-1
Enter fullscreen mode Exit fullscreen mode

Finally, I can get a token with the newly created user:

aws cognito-idp admin-initiate-auth --user-pool-id <pool_id> --client-id <client_id> --auth-flow ADMIN_USER_PASSWORD_AUTH --auth-parameters USERNAME=test_user@example.com,PASSWORD='Test1234!',SECRET_HASH=<hash> --region us-east-1
Enter fullscreen mode Exit fullscreen mode

Local test with Bruno

For development, it is most convenient to test the agent locally, without deploying every change to the cloud.

In the code, I rely on MEMORY_ID passed as an environment variable, so for local testing I create a .env file and I read it using dotenv. The Memory ID can be found in the .bedrock_agentcore.yaml

Since I am eventually going to use my agent from UI, I want to test it with HTTP request, instead of using agentcore invoke command.

I set headers x-amzn-bedrock-agentcore-runtime-custom-actorid for actor ID, and x-amzn-bedrock-agentcore-runtime-session-id for session ID.

To start local server: uv run main.py

And I can send a request. Looks good:

I can see in the logs that memory is loaded. I can test it locally, but I want to hit the live endpoint.

Testing live endpoint

To call our agent, we need to send a POST request to the endpoint:
https://bedrock-agentcore.<REGION>.amazonaws.com/runtimes/<URL_ENCODED_AGENT_ARN>/invocations?qualifier=DEFAULT

Agent ARN can be found in the .bedrock_agentcore.yaml file, or in the console.

I set the ActorId header as user_tech and I shared some insights with the agent:

Ok, it works fine as part of the same session_id. Maybe it doesn't look amazing, but consider that the AgentCore runtime is stateless.

Now let's try changing session_id

and ask the same question:

Great!

Let's now update session_id and actor_id to simulate a different user:

Nice :)

Summary

Using the AgentCore Memory service solves the context issues for stateless agent flows. With the session ID and actor ID, we are able to help the agent get relevant information and make the conversation more natural and less frustrating for the user.

Built-in strategies for storing memories are powerful. Strategies are also highly customizable. We can create our own logic for storing information.

AgentCore Memory is not a replacement for RAG. RAG is still relevant for getting domain-specific insights and general data. AgentCore Memory helps with keeping the context of the conversation related to the specific actor.

I didn't try other frameworks, but when using Strands, AgentCore Memory integration is seamless. I didn't need to set up any hooks and modify the code. It was enough to create a MemoryManager config and pass it to the agent.

Top comments (0)