Problem Statement

I would like Strands to allow me to pass in custom fields in a custom callback handler. Currently I only see the response token passed into the callback handler for response generation events. This is fine for printing out tokens for debugging but I would want to be able to do more with the callback handler.

Proposed Solution

Probably not the best way to do this, but there would be a chat session ID I would like to be able to access from the callback handler so that I can post response tokens to that chat session ID.

from strands import Agent
from strands_tools import calculator

def custom_callback_handler(**kwargs):
    # Process stream data
    if "data" in kwargs:
        print(f"MODEL OUTPUT: {kwargs['data']}")
        websocketAPI.post(kwargs["SessionID"], kwargs['data']['message'])
    elif "current_tool_use" in kwargs and kwargs["current_tool_use"].get("name"):
        print(f"\nUSING TOOL: {kwargs['current_tool_use']['name']}")

sessionID = message["SessionID"]
# Create an agent with custom callback handler
agent = Agent(
    tools=[calculator],
    callback_handler=custom_callback_handler(sessionID)
)

agent(message["Prompt"], sessionID)

Use Case

One use case is to be able to send response tokens to a particular chat session ID. The AWS AppSync Events service uses "channels" to route messages. (see the ID below) I'd like to be able to pass that ID into the callback handler so that each token can be routed to its appropriate channel.

payload = {
    'channel': f"/{OUTBOUND_NAMESPACE}/{OUTBOUND_CHANNEL}/{id}",
    'events': [json.dumps({'message': text})],
}
encoded_data = json.dumps(payload).encode('utf-8')
response = http.request('POST', API_ENDPOINT, body=encoded_data, headers=HEADERS)

Alternatives Solutions

No response

Additional Context

This could very well be a misunderstanding of how the callback handler works, my use case may be solvable with the callback handler as it exists today. Would appreciate an example of a callback handler with custom fields passed in, I wasn't able to get that to work.