Some lightweight cleanup and renaming for bedrock safety adapter

This commit is contained in:
Ashwin Bharambe 2024-09-24 19:27:03 -07:00
parent a2465f3f9c
commit f45705cd10
5 changed files with 76 additions and 78 deletions

View file

@ -4,21 +4,13 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Optional
from llama_models.schema_utils import json_schema_type
from pydantic import BaseModel, Field
import boto3
@json_schema_type
class BedrockShieldConfig(BaseModel):
class BedrockSafetyConfig(BaseModel):
"""Configuration information for a guardrail that you want to use in the request."""
aws_profile: Optional[str] = Field(
default='default',
aws_profile: str = Field(
default="default",
description="The profile on the machine having valid aws credentials. This will ensure separation of creation to invocation",
)