Strands Agent Hooks - BeforeModelnvocationEvent

[akingscote]

Hello,

Im working my way through issue196, and have integrated the Strands Agent SDK with Nvidia NeMo. I will get working examples with Guardrails AI and LLamaFirewall and will contribute to the strands examples, with complete documentation.

I'm using a the HookProvider with a BeforeModelInvocationEvent and the AfterInvocationEvent events.

However, I'm struggling to find an attribute in the event that cleanly contains the users inputs, and the outputs from LLMs.

I'm currently manually parsing event.agent.messages for the BeforeModelInvocationEvent, but it seems incredibly inefficient.

Can someone please advise if I am missing something?

Thanks, Ashley

Here is an example:

from strands.hooks import HookProvider, HookRegistry, AfterInvocationEvent
from strands.experimental.hooks import BeforeModelInvocationEvent
import json
from typing import Dict
import httpx

class CustomGuardrailHook(HookProvider):
    def register_hooks(self, registry: HookRegistry) -> None:
        registry.add_callback(BeforeModelInvocationEvent, self.check_input)
        # registry.add_callback(AfterInvocationEvent, self.check_output)

    def check_input(self, event: BeforeModelInvocationEvent) -> None:
        """
        We need to check the input for calls to the LLM, not just user-provided calls
        """
        # just yeet the entire message into the guardrail.
        message = extract_text_from_json(event.agent.messages[-1])
               
        payload= {
            "config_id": "nemo-guardrail-example",
            "messages": [{
                "role":"user",
                "content": message
            }]
        }

        headers = {
            "Content-Type": "application/json"
        }

        # Nvidia Nemo server mode
        url = "http://127.0.0.1:8000/v1/chat/completions"
        response = httpx.post(url, headers=headers, json=payload)

        # Raise an exception for bad status codes (4xx or 5xx)
        response.raise_for_status()

        if response.status_code != 200:
            raise Exception(f"Guardrail check failed with status code {response.status_code}")
    
        messages = response.json().get("messages")
        try:
            if messages[0].get("content") != "ALLOW":
                raise Exception("Guardrail check failed")

        except KeyError:
            raise Exception("Guardrail check failed: No messages returned from guardrail service")
        
        print("Guardrail check passed, proceeding with request.")

    def check_output(self, event: AfterInvocationEvent) -> None:
        print(f"Request completed for agent: {event.agent.name}")

def extract_text_from_json(message: Dict):
    """
    There has to be a better way to do this...
    """
    try:
        # Check if 'content' key exists and is a non-empty list
        if 'content' in message and message['content']:
            content_item = message['content'][0]

            # First, check for the 'toolResult' structure
            if 'toolResult' in content_item:
                # Navigate through the nested structure
                tool_content = content_item.get('toolResult', {}).get('content', [])
                if tool_content and 'text' in tool_content[0]:
                    return tool_content[0]['text']
            
            # If not, check for the direct 'text' key structure
            elif 'text' in content_item:
                return content_item['text']

        # Return None if the expected structure is not found
        return None

    except (json.JSONDecodeError, KeyError, IndexError, TypeError) as e:
        # Handle potential errors like invalid JSON, missing keys, or wrong types
        print(f"An error occurred: {e}")
        return None

CC @dbschmigelski

[akingscote]

@theagenticguy Do you have any ideas? cheers

[akingscote]

I think that the answer here is to just use the MessageAddedEvent hook, and send the entire message to any guardrails.

The event.agent.message objects have a pretty gnarly structure, and I don't know the schema. I figure the guardrail should be able to handle the message.

I'll leave this discussion open, but I think I have a better way forward

[zastrowm]

event.agent.messages would be the correct way; there's not a guarantee that the user's prompt was text (it could be an image or a document).

https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Message.html details the format of the Message object - it follows Bedrock conventions

[zastrowm]

There's probably a feature-request here to include the Message that will be added to the agent's messages on BeforeInvocationEvent