From 3eaa1fa217be9b6f176e46a2957a51ab408f1d5c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 29 Jul 2024 13:14:53 -0700
Subject: [PATCH] types add GuardrailConfigBlock

---
 litellm/types/llms/bedrock.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py
index 95ebc9742..cf0a4a84b 100644
--- a/litellm/types/llms/bedrock.py
+++ b/litellm/types/llms/bedrock.py
@@ -1,16 +1,18 @@
-from typing import TypedDict, Any, Union, Optional, Literal, List
 import json
-from .openai import ChatCompletionToolCallChunk
+from typing import Any, List, Literal, Optional, TypedDict, Union
+
 from typing_extensions import (
-    Self,
     Protocol,
-    TypeGuard,
-    override,
-    get_origin,
-    runtime_checkable,
     Required,
+    Self,
+    TypeGuard,
+    get_origin,
+    override,
+    runtime_checkable,
 )
 
+from .openai import ChatCompletionToolCallChunk
+
 
 class SystemContentBlock(TypedDict):
     text: str
@@ -108,6 +110,12 @@ class ToolConfigBlock(TypedDict, total=False):
     toolChoice: Union[str, ToolChoiceValuesBlock]
 
 
+class GuardrailConfigBlock(TypedDict, total=False):
+    guardrailIdentifier: str
+    guardrailVersion: str
+    trace: Literal["enabled", "disabled"]
+
+
 class InferenceConfig(TypedDict, total=False):
     maxTokens: int
     stopSequences: List[str]
@@ -144,6 +152,7 @@ class RequestObject(TypedDict, total=False):
     messages: Required[List[MessageBlock]]
     system: List[SystemContentBlock]
     toolConfig: ToolConfigBlock
+    guardrailConfig: Optional[GuardrailConfigBlock]
 
 
 class GenericStreamingChunk(TypedDict):