aws[patch]: add guardrails support

[baskaryan]

No description provided.

[3coins]

@supreetkt
Can you review this PR, and make sure to test the changes. Thanks.

[3coins]

@baskaryan
Thanks for making these changes. Some minor comments about the changes.

[3coins]

@baskaryan
Looks good. We should go ahead and merge.

One thing to note is that the usage from ChatBedrock is different which uses guardrails to set this config. We might want to align with the changes here and make guardrail_config as primary attribute, with guardrails as an alias there. But we can handle that in a separate PR. Thanks again for working on this.

[rsgrewal-aws]

hello Team, if the entry point will be ChatBedrock with converse api -- what will it mean for Guardrails config ? will it work as is ?

[3coins]

@rsgrewal-aws
I believe it will not work at the moment, as ChatBedrock is using a different field, and flattening the guardrail properties to pass to the invoke function. We will need few more changes on ChatBedrock to make it work. For example, this is currently failing in my test.

def test_guardrails() -> None:
    params = {
        "region_name": "us-west-2",
        "model_id": "anthropic.claude-3-sonnet-20240229-v1:0",
        "guardrails": {
            "guardrailIdentifier": "e7esbceow153",
            "guardrailVersion": "1",
            "trace": "enabled",
        },
        "beta_use_converse_api": True
    }
    chat_model = ChatBedrock(**params)  # type: ignore[arg-type]
    messages = [
        HumanMessage(
            content=[
                "Create a playlist of 2 heavy metal songs.",
                {
                    "guardContent": {
                        "text": {"text": "Only answer with a list of songs."}
                    }
                },
            ]
        )
    ]
    response = chat_model.invoke(messages)

    assert (
        response.content == "Sorry, I can't answer questions about heavy metal music."
    )
    assert response.response_metadata["stopReason"] == "guardrail_intervened"
    assert response.response_metadata["trace"] is not None

    stream = chat_model.stream(messages)
    response = next(stream)
    for chunk in stream:
        response += chunk

    assert (
        response.content[0]["text"]  # type: ignore[index]
        == "Sorry, I can't answer questions about heavy metal music."
    )
    assert response.response_metadata["stopReason"] == "guardrail_intervened"
    assert response.response_metadata["trace"] is not None

Error

E           botocore.errorfactory.ValidationException: An error occurred (ValidationException) when calling the Converse operation: The model returned the following errors: The guardrail can't assess the content in the guardContent field. The guardrail configuration is missing. Specify the guardrailConfig field and try again.

[3coins]

@baskaryan
Can you take a look at the recent commit, this should fix calling the ChatBedrock with guardrails config. I have added an integration test to verify this.

[supreetkt]

Maybe this can be a part of a separate PR, but would advise adding docstrings to all util methods. Otherwise LGTM.

[baskaryan]

is this specific to a user or will it work with our creds as well?