[Python] Invoking Bedrock Agent through Lambda reusing old inputs

I'm running into an issue where I'm trying to invoke a bedrock agent I've created through a lambda function within AWS. Testing the agent within the agent platform works fine, it's just invoking through my function, meant to be accessible through an external API call, that's causing me issues.For reference, I'm testing using Lambda's built in test cases, and all code is in python. Below is the code I've written for this:

import json
import boto3
import uuid

bedrock_agent_runtime = boto3.client('bedrock-agent-runtime')

AGENT_ID = 'xxxxx'
AGENT_ALIAS_ID = 'xxxxx'

def lambda_handler(event, context):
try:
body = json.loads(event['body'])
query = body.get('query', '')
data = body.get('data', '')

# Construct dynamic prompt
combined_input = f"""Use the following data to answer the question:\n{data}\n\nQuestion: {query}"""

session_id = str(uuid.uuid4())

print(f"Session ID: {session_id}")

response = bedrock_agent_runtime.invoke_agent(
agentId=AGENT_ID,
agentAliasId=AGENT_ALIAS_ID,
sessionId=session_id,
inputText=combined_input
)

stream = response['completion']
response_text = ""

for event in stream:
chunk = event.get('chunk')
if chunk and 'bytes' in chunk:
response_text += chunk['bytes'].decode('utf-8')

# This code is meant to double check the session is wiped.
bedrock_agent_runtime.invoke_agent(
agentId=AGENT_ID,
agentAliasId=AGENT_ALIAS_ID,
sessionId=session_id,
inputText= "Empty String",
endSession=True
)

return {
'statusCode': 200,
'headers': {"Content-Type": "application/json"},
'body': json.dumps({"response": response_text})
}

except Exception as e:
return {
'statusCode': 500,
'body': json.dumps({"error": str(e)})
}


This code works properly the first time in any given session. If I asked "What is the capitol of France?", it will provide the correct answer. The problem starts with the second question. If I ask something like "What is the capitol of Russia", the agent is somehow receiving the first prompt about France again and generating an answer for that. If I ask a third question about the capitol of China, then it will receive the question about Russia and answer it, and so forth, always one prompt behind.

I can't figure out what could be causing this. I believe I've set a unique ID to be created for each query with session_id = str(uuid.uuid4()), so each question should receive a new session, and I've also tried explicitly ending the old session with the second call to the agent api right before my return statement, neither of which are having any effect. What could be the issue here? Again, this doesn't happen when testing the agent through Bedrock directly.

Continue reading...